Numpy:将方阵的索引转换为上三角索引的有效方法

Question

问题:给定一个索引元组,返回上三角指数的顺序.这是一个例子:

假设我们有一个形状的方阵A(3,3).

A具有6个上三角形索引,即(0,0),(0,1),(0,2),(1,1),(1,2),(2,2).

现在我知道索引(1,2)处的元素,它是属于A的上三角部分的索引.我想返回4(这意味着它是所有上三角索引中的第5个元素.)

关于如何做到这一点的任何想法？

最好的,志浩

Answer 1

人们可以写出明确的公式:

def utr_idx(N, i, j):
    return (2*N+1-i)*i//2 + j-i

演示:

>>> N = 127
>>> X = np.transpose(np.triu_indices(N))
>>> utr_idx(N, *X[2123])
2123