如何在numpy ndarray中获得每行的N个最大值?

时间:2016-05-09 21:08:48

标签: python numpy

当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>时,获得前N的最佳方法是什么? 1? (比方说,5)

3 个答案:

答案 0 :(得分:3)

使用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]

如果您的数组足够大,这仍然比整个数组上的完整排序快得多。

要对二维数组的每一行进行排序,您可以使用axis=和/或np.partition的{​​{1}}个参数:

np.sort

基准:

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]]

答案 1 :(得分:2)

为什么不这样做:

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

答案 2 :(得分:2)

partitionsortargsort等采用轴参数

让我们改变一些值

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个; slice获取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]])