使用 np.argpartition
和 np.argsort
两种方法对 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
和k_largest_index_argpartition_v2
之间的差异在于我们如何使用argparition
。在第一个版本中,我们对输入数组取负值,然后使用argpartition
获取最小的k
个索引,从而有效地获取最大的k
个索引;而在第二个版本中,我们获取前a.size-k
个最小的索引,然后选择剩余的最大的k
个索引。
此外,在使用argpartition
时,值得注意的是我们不能按照它们的排序顺序获取索引。如果需要按排序顺序获取索引,则需要将范围数组提供给np.argpartition
,正如这个post
中所述。
样例运行 -
1) 二维情况:
In [42]: a
Out[42]:
array([[38, 14, 81, 50],
[17, 65, 60, 24],
[64, 73, 25, 95]])
In [43]: k_largest_index_argsort(a, k=2)
Out[43]:
array([[2, 3],
[0, 2]])
In [44]: k_largest_index_argsort(a, k=4)
Out[44]:
array([[2, 3],
[0, 2],
[2, 1],
[1, 1]])
In [66]: k_largest_index_argpartition_v1(a, k=4)
Out[66]:
array([[2, 1],
[2, 3],
[0, 2],
[1, 1]])
2) 三维情况:
In [46]: a
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