时间戳

首页 > 解决方案 > 在一个 numpy 2D 数组上 - 当 N 更改行时如何将每行中的最后 N 个数组元素设置为零

python - 在一个 numpy 2D 数组上 - 当 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]])

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

标签： pythonnumpyvectorization

解决方案

用于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方法是 with 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]])

推荐阅读