python之在 numpy 二维数组上之当 N 改变行时如何将每行中的最后 N 个数组元素设置为零

比如说，我有一个由 20 个元素组成的 2D numpy 数组，例如:

arr = np.array([[1, 2, 15, 7],[9, 11, 17, 19],[5, 7, 5, 8],[19, 4, 1, 45],[10, 7, 14, 8]]) 

和一个额外的数组:

to_zero = np.array([0, 2, 1, 3, 2]) 

现在，对于每一行 i，我想让最后的 to_zero[i] 元素等于零，所以最终我们会得到以下结果:

res = np.array([[1, 2, 15, 7],[9, 11, 0, 0],[5, 7, 5, 0],[19, 0, 0, 0],[10, 7, 0, 0]]) 

我想在一个非常大的数组上做这个操作。 有什么方法可以将此操作矢量化，没有循环，也没有辅助数组吗？

请您参考如下方法:

使用 broadcasted-comparison 获取那些尾部的掩码，然后掩码输入 -

In [63]: r = np.arange(arr.shape[1])[::-1] 
 
In [66]: mask = to_zero[:,None]>r 
 
In [69]: mask # mask of trailing places to be reset in input 
Out[69]:  
array([[False, False, False, False], 
       [False, False,  True,  True], 
       [False, False, False,  True], 
       [False,  True,  True,  True], 
       [False, False,  True,  True]]) 
 
In [67]: arr[mask] = 0 
 
In [68]: arr 
Out[68]:  
array([[ 1,  2, 15,  7], 
       [ 9, 11,  0,  0], 
       [ 5,  7,  5,  0], 
       [19,  0,  0,  0], 
       [10,  7,  0,  0]]) 

获取 r 的替代方法是使用 np.arange(arr.shape[1]-1,-1,-1)。

使用逐元素乘法获得最终输出的替代方法:arr*~mask。

或者用翻转比较构造翻转掩码然后相乘-

In [75]: arr*(to_zero[:,None]<=np.arange(arr.shape[1]-1,-1,-1)) 
Out[75]:  
array([[ 1,  2, 15,  7], 
       [ 9, 11,  0,  0], 
       [ 5,  7,  5,  0], 
       [19,  0,  0,  0], 
       [10,  7,  0,  0]]) 

对于大型阵列，利用 numexpr 的多核 -

In [78]: import numexpr as ne 
 
In [79]: ne.evaluate('arr*mask',{'mask':to_zero[:,None]<=np.arange(arr.shape[1]-1,-1,-1)}) 
Out[79]:  
array([[ 1,  2, 15,  7], 
       [ 9, 11,  0,  0], 
       [ 5,  7,  5,  0], 
       [19,  0,  0,  0], 
       [10,  7,  0,  0]])