m Smallest values from upper triangular matrix with their indices as a list of tuples

Question

I have a np.ndarray as follows:

[[ inf   1.   3.   2.   1.]
 [ inf  inf   2.   3.   2.]
 [ inf  inf  inf   5.   4.]
 [ inf  inf  inf  inf   1.]
 [ inf  inf  inf  inf  inf]]

Is there a way to get the indices and values of the m smallest items in that nd array? So, if I wanted the 4 smallest it would be

[(0,1,1),(0,4,1),(3,4,1),(0,3,2)] 

where (row,col,val) is the notation above.

If there are multiple values then one of them is just randomly chosen. For instance, there were 3 ones and then next smallest is a value 2 but (0,3,2), (1,2,2),(1,4,2) were all possible choices.

Essentially, Can I extract the k smallest values in that format from the upper triangular matrix efficiently (the matrix is much larger than the example above). I tried flattening it, using square form, nsmallest, but am having trouble getting the indices and values to align. Thanks!

Answer 1

For an Inf filled array -

r,c = np.unravel_index(a.ravel().argsort()[:4], a.shape)
out = zip(r,c,a[r,c])

For performance, consider using np.argpartition. So, replace a.ravel().argsort()[:4] with np.argpartition(a.ravel(), range(4))[:4].

Sample run -

In [285]: a
Out[285]: 
array([[ inf,   1.,   3.,   2.,   1.],
       [ inf,  inf,   2.,   3.,   2.],
       [ inf,  inf,  inf,   5.,   4.],
       [ inf,  inf,  inf,  inf,   1.],
       [ inf,  inf,  inf,  inf,  inf]])

In [286]: out
Out[286]: [(0, 1, 1.0), (0, 4, 1.0), (3, 4, 1.0), (0, 3, 2.0)]

---

For a generic case -

R,C = np.triu_indices(a.shape[1],1)
idx = a[R,C].argsort()[:4]
r,c = R[idx], C[idx]
out = zip(r,c,a[r,c])

Sample run -

In [351]: a
Out[351]: 
array([[ 68.,  67.,  81.,  23.,  16.],
       [ 84.,  83.,  20.,  66.,  48.],
       [ 58.,  72.,  98.,  63.,  30.],
       [ 61.,  40.,   1.,  86.,  22.],
       [ 29.,  95.,  38.,  22.,  95.]])
In [352]: out
Out[352]: [(0, 4, 16.0), (1, 2, 20.0), (3, 4, 22.0), (0, 3, 23.0)]

For performance, consider using np.argpartition. So, replace a[R,C].argsort()[:4] with np.argpartition(a[R,C], range(4))[:4].