Numpy：将正方形矩阵的索引转换为其上三角索引的有效方法

Question

给定一个索引元组，返回它在上三角索引中的顺序。以下是一个例子：

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

A有6个上三角索引，分别为(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)。

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

你有什么通用的解决方法吗？

最好, Zhihao

- Zhihao Cui

---

更一般地说，如果我有一个索引列表，是否有一个函数可以将它们一起转换为triu索引[返回一个转换结果的列表/数组]? - Zhihao Cui

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

- Paul Panzer

---

我刚刚意识到了，哈哈 - Zhihao Cui

逆过程怎么样？从向量返回矩阵（i，j）索引？ - seralouk

Willem Van Onsem在他的回答@makis中写下了反函数。 - Paul Panzer

Answer 2

对于一个 n×n 的矩阵，上三角的 (i, j) 元素是矩阵的第 i×(2×n-i+1)/2+j-i 个元素。

我们也可以反过来计算，通过以下公式得到第 k 个元素的 (i, j) 坐标： i = ⌊(-√((2n+1)²-8k)+2n+1)/2⌋ 和 j = k+i-i×(2×n-i+1)/2 例如：

from math import floor, sqrt

def coor_to_idx(n, i, j):
    return i*(2*n-i+1)//2+j-i

def idx_to_coor(n, k):
    i = floor((-sqrt((2*n+1)*(2*n+1)-8*k)+2*n+1)/2)
    j = k + i - i*(2*n-i+1)//2
    return i, j

例如：

>>> [idx_to_coor(4, i) for i in range(10)]
[(0, 0), (0, 1), (0, 2), (0, 3), (1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]
>>> [coor_to_idx(4, i, j) for i in range(4) for j in range(i, 4)]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

考虑到数值并不巨大（如果这些数值很大，计算就不再是常数时间），因此我们可以在O(1)的时间内计算第k个坐标，例如：

>>> idx_to_coor(1234567, 123456789)
(100, 5139)

相当于通过枚举获得它：

>>> next(islice(((i, j) for i in range(1234567) for j in range(i, 1234567)), 123456789, None))
(100, 5139)

在这里，将索引转换为坐标也可能会由于浮点精度问题而产生一些舍入误差，特别是对于大数值。

- Willem Van Onsem

Answer 3

如果我理解得正确，您可以使用 itertools 的组合函数（带有替换）来获取索引。

>>> ind = tuple(itertools.combinations_with_replacement(range(3),2))
((0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2))

要检索索引，只需使用index方法。

>>> ind.index((1,2))
4

- rafaelc

Answer 4

你可以使用 np.triu_indices 和一个 字典：

import numpy as np

iu1 = np.triu_indices(3)
table = {(i, j): c for c, (i, j) in enumerate(zip(*iu1))}
print(table[(1, 2)])

输出

4

- Dani Mesejo

Answer 5

与 @DanielMesejo 相似，您可以使用 np.triu_indices 与 argwhere 或 nonzero：

my_index = (1,2)

>>> np.nonzero((np.stack(np.triu_indices(3), axis=1) == my_index).all(1))
(array([4]),)
>>> np.argwhere((np.stack(np.triu_indices(3), axis=1) == my_index).all(1))
array([[4]])

说明：

np.stack(np.triu_indices(3), axis=1)会按顺序给出上三角的索引：

array([[0, 0],
       [0, 1],
       [0, 2],
       [1, 1],
       [1, 2],
       [2, 2]])

所以，你需要做的就是找到它与[1,2]匹配的地方（你可以使用==操作符和all实现）

- sacuL

Answer 6

构建上标会很耗费时间。我们可以直接这样获取相应的索引 -

def triu_index(N, x, y):
    # Get index corresponding to (x,y) in upper triangular list
    idx = np.r_[0,np.arange(N,1,-1).cumsum()]
    return idx[x]+y-x

示例运行 -

In [271]: triu_index(N=3, x=1, y=2)
Out[271]: 4

- Divakar