Numpy: Efficient way to convert indices of a square matrix to its upper triangular indices [closed]

Question

Question: given a tuple of index, return its order in upper triangular indices. Here is an example:

Suppose we have a square matrix A of shape (3, 3).

A has 6 upper triangular indices, namely, (0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2).

Now I know an element at index (1, 2), which is a index belongs to the upper triangular part of A. I would like to return 4 (which means it is the 5th element in all upper triangular indices.)

Any ideas on how to do that in general?

Best, Zhihao

Answer 1

One can write down the explicit formula:

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

Demo:

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