python 数组取前k最大,如何从多维数组中获取k个最大值的索引

使用^{}和{a2}对ndarrays-def k_largest_index_argpartition_v1(a, k):

idx = np.argpartition(-a.ravel(),k)[:k]

return np.column_stack(np.unravel_index(idx, a.shape))

def k_largest_index_argpartition_v2(a, k):

idx = np.argpartition(a.ravel(),a.size-k)[-k:]

return np.column_stack(np.unravel_index(idx, a.shape))

def k_largest_index_argsort(a, k):

idx = np.argsort(a.ravel())[:-k-1:-1]

return np.column_stack(np.unravel_index(idx, a.shape))

用argpartition讨论两个版本

k_largest_index_argpartition_v1和{}之间的区别在于我们如何使用argparition。在第一个版本中，我们对输入数组求反，然后使用argpartition来获得最小的k索引的索引，从而有效地得到最大的索引，而在第二个版本中，我们得到了第一个最小的索引，然后选择剩下的最大的索引。在

另外，这里值得一提的是，使用argpartition，我们没有按排序顺序得到索引。如果需要排序顺序，我们需要向np.argpartition提供范围数组，如^{}所述。在

样本运行-

1)二维盒：

^{pr2}$

2)3D案例：In [46]: a # 3D array

Out[46]:

array([[[20, 98, 27, 73],

[33, 78, 48, 59],

[28, 91, 64, 70]],

[[47, 34, 51, 19],

[73, 38, 63, 94],

[95, 25, 93, 64]]])

In [47]: k_largest_index_argsort(a, k=2)

Out[47]:

array([[0, 0, 1],

[1, 2, 0]])

运行时测试-In [56]: a = np.random.randint(0,99999999999999,(3000,4000))

In [57]: %timeit k_largest_index_argsort(a, k=10)

1 loops, best of 3: 2.18 s per loop

In [58]: %timeit k_largest_index_argpartition_v1(a, k=10)

10 loops, best of 3: 178 ms per loop

In [59]: %timeit k_largest_index_argpartition_v2(a, k=10)

10 loops, best of 3: 128 ms per loop