查找 2D 数组/矩阵中 k 个最高值的索引

2024-02-23

我有一个包含值的 2D 矩阵,我想找到前 5 个值的索引。 例如对于

matrix([[0.17542851, 0.13199346, 0.01579704, 0.01429822, 0.01302919],
        [0.13279703, 0.12444886, 0.04742024, 0.03114371, 0.02623729],
        [0.13502306, 0.07815065, 0.07291175, 0.03690815, 0.02163695],
        [0.19032505, 0.15853737, 0.05889324, 0.02791679, 0.02699252],
        [0.1695696 , 0.14538635, 0.07127667, 0.04997876, 0.02580234]])

我想要得到(0,3), (0,1), (0,4), (3,1), (4,1)

我搜索并尝试了很多解决方法,包括np.argmax(), np.argsort(), np.argpartition()没有任何好的结果。 例如:

>>np.dstack(np.unravel_index(np.argsort(a.ravel(),axis=None), a.shape))

array([[[0, 4],
        [0, 3],
        [0, 2],
        [2, 4],
        [4, 4],
        [1, 4],
        [3, 4],
        [3, 3],
        [1, 3],
        [2, 3],
        [1, 2],
        [4, 3],
        [3, 2],
        [4, 2],
        [2, 2],
        [2, 1],
        [1, 1],
        [0, 1],
        [1, 0],
        [2, 0],
        [4, 1],
        [3, 1],
        [4, 0],
        [0, 0],
        [3, 0]]], dtype=int64)

这个结果没有任何意义。 请注意,我想要原始索引,我不关心顺序(只想要任意顺序的前 5 个索引,但升序会更好)


np.argpartition https://docs.scipy.org/doc/numpy/reference/generated/numpy.argpartition.html应该是一个很好的工具(高效的工具)来获得这些顶部k索引而不维持秩序。因此,对于数组数据a, 这将是 -

In [43]: np.c_[np.unravel_index(np.argpartition(a.ravel(),-5)[-5:],a.shape)]
Out[43]: 
array([[4, 1],
       [3, 1],
       [4, 0],
       [0, 0],
       [3, 0]])

为了解释一下,让我们将其分解为单个流程步骤 -

# Get partitioned indices such that the last 5 indices refer to the top 5
# values taken globally from the input array. Refer to docs for more info
# Note that it's global because we will flatten it. 
In [9]: np.argpartition(a.ravel(),-5)
Out[9]: 
array([14, 24,  2,  3,  4, 23, 22,  7,  8,  9, 19, 18, 17, 13, 12, 11,  6,
        1,  5, 10, 21, 16, 20,  0, 15])

# Get last 5 indices, which are the top 5 valued indices
In [10]: np.argpartition(a.ravel(),-5)[-5:]
Out[10]: array([21, 16, 20,  0, 15])

# Convert the global indices back to row-col format
In [11]: np.unravel_index(np.argpartition(a.ravel(),-5)[-5:],a.shape)
Out[11]: (array([4, 3, 4, 0, 3]), array([1, 1, 0, 0, 0]))

# Stack into two-columnar array
In [12]: np.c_[np.unravel_index(np.argpartition(a.ravel(),-5)[-5:],a.shape)]
Out[12]: 
array([[4, 1],
       [3, 1],
       [4, 0],
       [0, 0],
       [3, 0]])

对于矩阵数据a, 这将是 -

In [48]: np.dstack(np.unravel_index(np.argpartition(a.ravel(),-5)[:,-5:],a.shape))
Out[48]: 
array([[[4, 1],
        [3, 1],
        [4, 0],
        [0, 0],
        [3, 0]]])

因此,与数组相比,唯一的区别是使用np.dstack,因为对于矩阵数据,数据始终保持二维状态。

请注意,这些是您上次的结果5 rows.

本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)

查找 2D 数组/矩阵中 k 个最高值的索引 的相关文章

随机推荐