使用np.triu_indices生成索引

Question

我想在上三角矩阵的每个条目上调用一个函数。为了避免for循环的混乱嵌套，我使用的是numpy.triu_indices函数。我的功能有效，但我想知道是否有更简洁的方法来实现索引。

import numpy as np
def weighted_edges(adjmat):
    indices = np.triu_indices(len(adjmat))
    return ((x, y, adjmat[x,y]) for (x,y) in zip(indices[0], indices[1]))

我怀疑有一种方法可以实现这一点而无需在zip调用中引用indices [i]。确实有办法这样做吗？

Answer 1

如果你有一个N x N矩阵，你想要上三角形值，只需做

import numpy as np
N = 5
x = np.arange(N**2).reshape(N, N)
upper = x[np.triu_indices(N, 0)]

如果您希望三角形值偏离主对角线k列，请执行

upper = x[np.triu_indices(N, k)]