python - m 上三角矩阵中的最小值，其索引为元组列表

Question

我有一个 np.ndarray 如下:

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

有没有办法获取那个nd数组中m个最小项的索引和值？所以，如果我想要 4 个最小的，它将是

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

其中 (row,col,val) 是上面的符号。

如果有多个值，则随机选择其中一个。例如，有 3 个，然后下一个最小的是值 2，但 (0,3,2)、(1,2,2)、(1,4,2) 都是可能的选择。

本质上，我能否有效地从上三角矩阵中提取该格式的 k 个最小值(该矩阵比上面的示例大得多)。我尝试使用正方形 nsmallest 将其展平，但无法使索引和值对齐。谢谢!

Answer 1

对于 Inf 填充数组 -

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

为了提高性能，请考虑使用 np.argpartition。因此，将 a.ravel().argsort()[:4] 替换为 np.argpartition(a.ravel(), range(4))[:4] .

sample 运行-

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

---

对于一般情况 -

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 运行-

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

为了提高性能，请考虑使用 np.argpartition。因此，将 a[R,C].argsort()[:4] 替换为 np.argpartition(a[R,C], range(4))[:4].