python - 使用 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)]