说,我有一个由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]])
我想在很大的阵列上执行此操作。 有没有办法向量化此操作,没有循环,也没有辅助数组?
答案 0 :(得分:5)
使用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
。
或用 flipped-comparison 构造 flipped-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]])