python - 保持原始数组维度的二维数组的 argpartsort/partsort

Question

我有一个大型二维数组(例如 [1000, 100])，我需要对其进行元素分类。我需要获取数组每一行中前 n 个最大的项，但我需要将所有项保留在它们的位置并将所有其他项替换为 0。

例如对于 3x5 数组的每行前 3 项:

input:
   [[1, 2, 3, 4, 5],
    [5, 4, 3, 2, 1],
    [8, 5, 1, 9, 4]]

output:
   [[0, 0, 3, 4, 5],
    [5, 4, 3, 0, 0],
    [8, 5, 0, 9, 0]]

我可以通过执行 bottleneck.partsort 来慢慢地做到这一点，以获取每行的前 n 个项目，然后将第 n 个值以下的所有内容设置为零:

for row in input:
    row[row < -partsort(-row, 3)[:3][-1]] = 0

有没有更快的方法来做到这一点？对于有很多行的较大阵列，这可能需要一段时间...

谢谢

Answer 1

你可以做-

n = 3         # Number of elements to keep per row
A[np.arange(A.shape[0])[:,None],A.argsort(1)[:,:A.shape[1]-n]] = 0

sample 运行-

In [38]: A
Out[38]: 
array([[ 1, 85, 59,  1, 67, 33,  6, 61],
       [ 5, 81, 72, 14, 43, 76, 23, 23],
       [67, 49, 76, 22, 58, 66,  5, 74],
       [40, 68, 11, 65, 55, 58, 49, 73]])

In [39]: n = 3 # Number of elements to keep per row
    ...: A[np.arange(A.shape[0])[:,None],A.argsort(1)[:,:A.shape[1]-n]] = 0
    ...: 

In [40]: A
Out[40]: 
array([[ 0, 85,  0,  0, 67,  0,  0, 61],
       [ 0, 81, 72,  0,  0, 76,  0,  0],
       [67,  0, 76,  0,  0,  0,  0, 74],
       [ 0, 68,  0, 65,  0,  0,  0, 73]])