如何在numpy ndarray中每行获取N个最大值？

Question

我们知道当N=1时该如何做。

import numpy as np

m = np.arange(15).reshape(3, 5)
m[xrange(len(m)), m.argmax(axis=1)]    # array([ 4,  9, 14])

当N>1时（比如说，5），获取前N个元素的最佳方式是什么？

- Jay

Answer 1

使用np.partition进行部分排序比完全排序要便宜得多。

gen = np.random.RandomState(0)
x = gen.permutation(100)

# full sort
print(np.sort(x)[-10:])
# [90 91 92 93 94 95 96 97 98 99]

# partial sort such that the largest 10 items are in the last 10 indices
print(np.partition(x, -10)[-10:])
# [90 91 93 92 94 96 98 95 97 99]

如果您需要排序最大的N个项目，可以在您部分排序的数组中对最后N个元素调用np.sort：

print(np.sort(np.partition(x, -10)[-10:]))
# [90 91 92 93 94 95 96 97 98 99]

如果您的数组足够大，使用 np.partition 和/或 np.sort 的 axis= 参数可以在二维数组的每一行中进行排序，这仍然比整个数组的完全排序要快得多。

---

y = np.repeat(np.arange(100)[None, :], 5, 0)
gen.shuffle(y.T)

# partial sort, followed by a full sort of the last 10 elements in each row
print(np.sort(np.partition(y, -10, axis=1)[:, -10:], axis=1))
# [[90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]
#  [90 91 92 93 94 95 96 97 98 99]]

---

基准测试：

In [1]: %%timeit x = np.random.permutation(10000000)
   ...: np.sort(x)[-10:]
   ...: 
1 loop, best of 3: 958 ms per loop

In [2]: %%timeit x = np.random.permutation(10000000)
np.partition(x, -10)[-10:]
   ....: 
10 loops, best of 3: 41.3 ms per loop

In [3]: %%timeit x = np.random.permutation(10000000)
np.sort(np.partition(x, -10)[-10:])
   ....: 
10 loops, best of 3: 78.8 ms per loop

- ali_m

---

如果您确实需要对它们进行排序，使用np.partition然后对您关心的部分进行排序仍然可以比完全排序更高效。 - user2357112

@user2357112，非常正确，我已经更新了我的答案，包括基准测试。 - ali_m

Answer 2

为什么不像这样做呢：

np.sort(m)[:,-N:]

- Jules

Answer 3

partition、sort、argsort等函数都需要一个轴参数。

让我们对一些值进行洗牌。

In [161]: A=np.arange(24)

In [162]: np.random.shuffle(A)

In [163]: A=A.reshape(4,6)

In [164]: A
Out[164]: 
array([[ 1,  2,  4, 19, 12, 11],
       [20,  5, 13, 21, 22,  3],
       [10,  6, 16, 18, 17,  8],
       [23,  9,  7,  0, 14, 15]])

分区：

In [165]: A.partition(4,axis=1)

In [166]: A
Out[166]: 
array([[ 2,  1,  4, 11, 12, 19],
       [ 5,  3, 13, 20, 21, 22],
       [ 6,  8, 10, 16, 17, 18],
       [14,  7,  9,  0, 15, 23]])

每行前4个最小值排在前面，后2个最大值排在最后；切片获取包含最大值的数组：

In [167]: A[:,-2:]
Out[167]: 
array([[12, 19],
       [21, 22],
       [17, 18],
       [15, 23]])

排序可能较慢，但在这样一个小的数组上可能并不重要。此外，它可以让您选择任何N。

In [169]: A.sort(axis=1)

In [170]: A
Out[170]: 
array([[ 1,  2,  4, 11, 12, 19],
       [ 3,  5, 13, 20, 21, 22],
       [ 6,  8, 10, 16, 17, 18],
       [ 0,  7,  9, 14, 15, 23]])

- hpaulj