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

user_9987089 0
Viewed 1134

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

matrix numpy python
6 Answers

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
user_72073920

For an n×n matrix, the (i, j)-th item of the upper triangle is the i×(2×n-i+1)/2+j-i-th element of the matrix.

We can also do the math in reverse and obtain the (i, j) element for the k-th element with:

i = ⌊(-√((2n+1)2-8k)+2n+1)/2⌋ and j = k+i-i×(2×n-i+1)/2

So for example:

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

For example:

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

Given the numbers are not huge (well if these are huge, calculations are no longer done in constant time), we can thus calculate the k-th coordinate in O(1), for example:

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

which is equivalent to obtaining it through enumeration:

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

Here converting indices to a coordinate can also have, for large numbers, some rounding errors due to floating point imprecision.

user_675790

IIUC, you can get the indexes using itertools combinations with replacement 

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

To retrieve the index, just use index method 

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

You could use np.triu_indices and a dictionary:

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

Output

4
user_40015920

Similar to @DanielMesejo, you can use np.triu_indices with either argwhere or 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]])

Explanation:

np.stack(np.triu_indices(3), axis=1) gives you the indices of your upper triangle in order:

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

So all you have to do is find where it matches [1,2] (which you can do with the == operator and all)

user_66711760

Constructing upper indices would be costly. We can directly get the corresponding index like so -

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

Sample run -

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