我有一个带有值的二维矩阵，我想找到前 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应该是一个很好的工具(有效的工具)，可以在不维护秩序的情况下获得这些前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的使用 ，因为对于矩阵数据，数据始终保持为 2D。
请注意，这些是最后5行的结果。
您的示例：
n = np.array([[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]])
您的输出不是前 5 个值的标记。前 5 个值是
array([0.14538635, 0.15853737, 0.1695696 , 0.17542851, 0.19032505])
要获取他们的索引：sort并使用isin标记他们的位置True。最后，使用argwhere来获得他们的posision
np.argwhere(np.isin(n, np.sort(n, axis=None)[-5:]))
Out[324]:
array([[0, 0],
[3, 0],
[3, 1],
[4, 0],
[4, 1]], dtype=int32)
假设你有一个列表列表：
In [112]: M
Out[112]:
[[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]]
In [113]: heapq.nlargest(5, ((r,c) for r in range(len(M)) for c in range(len(M[r]))), key=lambda t: M[t[0]][t[1]])
Out[113]: [(4, 4), (4, 3), (4, 2), (4, 1), (4, 0)]
别忘了import heapq
我从一个引用@Divakar(非常优雅和快速)答案的问题来到这里。
排名的一个常见问题是如何处理重复(并列)。
在某些情况下，最好使用"密集排名"，其中[4, 7, 7, 9]将排名(按升序排列)：[0, 1, 1, 2].
相比之下，@Divakar的回答基本上反映了"序数排名"，其中[4, 7, 7, 9]将按升序排列[0, 1, 2, 3]。 在"topk"问题中，这可能有点违反直觉。例如，在：
b = np.array(
[[8, 6, 3],
[6, 7, 2],
[0, 8, 9]])
随着等级k=2和(并假设降序)，它给出：
k = 2
>>> np.c_[np.unravel_index(np.argpartition(b.ravel(),-k)[-k:], b.shape)]
array([[2, 1],
[2, 2]])
对应于9，并且仅对应于8值中的一个，而忽略了另一个8值。
如果有人对"密集排名"感兴趣，我会提出以下建议(它以"任何顺序"返回前 k 个值的所有索引 - 实际上，按索引顺序)：
def topk_indices(a, k):
_, rix = np.unique(-a, return_inverse=True)
return np.c_[np.unravel_index(np.where(rix < k)[0], a.shape)]
在 OP 的数组上：
>>> topk_indices(a, 5)
array([[0, 0],
[3, 0],
[3, 1],
[4, 0],
[4, 1]])
在上面的 int 数组上：
>>> topk_indices(b, 2)
array([[0, 0],
[2, 1],
[2, 2]])
性能
在性能方面，@Divakar的答案比这快 5 到 10 倍，适用于不同尺寸和参数的广泛测试。所以如果你认为你没有关系，或者你不在乎，那就用它代替这个。
举个例子：
a = np.random.randint(0, 10, (1_000_000, 2))
t0 = %timeit -o topk_indices(a, 5)
# 157 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
t1 = %timeit -o divakar_topk_indices(a, 5)
# 25.1 ms ± 49.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> t0.average / t1.average
6.24
作为旁注，它冒犯了我的敏感性，必须对整个数组(O(n log n))进行排序才能找到top-k...更合乎逻辑的heapq方法扩展得更好(O(n log k))，但具有更大的常量乘数(仅heapq.nlargest(5, a.ravel())需要 211 毫秒，这只是返回值，而不是索引。