我有一个2维numpy数组:
arr = np.array([[0.1, 0.1, 0.3, 0.4, 0.5],
[0.06, 0.1, 0.1, 0.1, 0.01],
[0.24, 0.24, 0.24, 0.24, 0.24],
[0.2, 0.25, 0.3, 0.12, 0.02]])
print (arr)
[[0.1 0.1 0.3 0.4 0.5 ]
[0.06 0.1 0.1 0.1 0.01]
[0.24 0.24 0.24 0.24 0.24]
[0.2 0.25 0.3 0.12 0.02]]
我想筛选出前N个值,所以我使用了 argsort
:
N = 2
arr1 = np.argsort(-arr, kind='mergesort') < N
print (arr1)
[[False False False True True]
[ True False False True False] <- first top 2 are duplicates
[ True True False False False]
[False True True False False]]
它表现良好,至少不会有像第二行那样的完全重复。
期望输出:
print (arr1)
[[False False False True True]
[False True True False False]
[ True True False False False]
[False True True False False]]
有没有更快的方法来处理它?
np.argsort(-arr, kind='mergesort')[:,:N]
吗? - Divakar