当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)
答案 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)
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个; 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]])