Question

我有一个大的2D数组（例如[1000,100]），我需要做一个元素方式的部分。我需要在数组的每一行中获得前n个最大的项目，但我需要将所有项目保留在其位置，并将所有其他条目替换为0。

E.g。对于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

示例运行 -

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

Answer 2

您可以在axis=1上使用np.partition并让broadcasting执行for循环：

>>> a
array([[1, 2, 3, 4, 5],
       [5, 4, 3, 2, 1],
       [8, 5, 1, 9, 4]])

>>> w = np.partition(a, -3, axis=1)[:, -3]
>>> a[a < w[:, np.newaxis]] = 0
>>> a
array([[0, 0, 3, 4, 5],
       [5, 4, 3, 0, 0],
       [8, 5, 0, 9, 0]])

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

2 个答案: