将 numpy 数组轴归零的 Pythonic 方法

Question

有点恼火，我自己没能得到这个，所以就这样吧。

假设我有一个二维 numpy 数组

import numpy as np
a = np.array([[1,2,3],[4,5,6],[7,8,9]])

我现在的情况是，我想将此数组中的一列清零，但保留原始数组。我可以通过以下方式做到这一点

b = a.copy()
b[:,1] = 0

获取

array([[1, 0, 3],
       [4, 0, 6],
       [7, 0, 9]])

或者如果我想将一系列列归零，我可以这样做

b = a.copy()
b[:,[0,3]] = 0

获取

array([[0, 2, 0],
       [0, 5, 0],
       [0, 8, 0]])

b 数组只会使用一次，然后就被丢弃。

是否有更多的 Pythonic 方法可以作为一个单行来完成此操作？我只希望零值 b 数组传递绘图例程，之后就不需要了。本质上，我不想在调用我的绘图函数之前有额外的两行 - 如果我可以在调用我的例程时做到这一点，那么它会干净得多。例如

plotting_func(<numpy_magic_here>)

而不是

b = a.copy()
b[:,1] = 0
plotting_func(b)

仅让 b 不再被使用

Answer 1

要将某些列设置为零，您可以使用 np.in1d 和 np.arange 来创建无效元素的掩码，该掩码与输入数组相乘时将使用 NumPy 广播将无效列设置为零。因此，我们将有一个单行实现，如下所示 -

a*~np.in1d(np.arange(a.shape[1]),cols_to_be_reset) 

或者，可以使用 np.where 来选择而不是相乘，就像这样 -

np.where(np.in1d(np.arange(a.shape[1]),cols_to_be_reset),0,a)

请注意，这并不是为了性能，而只是作为一句台词。

示例运行 -

In [546]: a
Out[546]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [547]: a*~np.in1d(np.arange(a.shape[1]),1)
Out[547]: 
array([[1, 0, 3],
       [4, 0, 6],
       [7, 0, 9]])

In [548]: a*~np.in1d(np.arange(a.shape[1]),[0,2])
Out[548]: 
array([[0, 2, 0],
       [0, 5, 0],
       [0, 8, 0]])