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)]