python – 根据布尔向量选择numpy中的列

我有两个NumPy数组a,b,维度为m乘n.我有一个长度为n的布尔向量b,我想生成一个新的数组c,它从a,b中选择n列,这样如果b [i]为真,我从b获取列,否则从a获取.

我如何以最有效的方式做到这一点？
我看了选择,在哪里选择.

解决方法

首先,让我们设置一些示例代码：

import numpy as np

m,n = 5,3
a = np.zeros((m,n))
b = np.ones((m,n))

boolvec = np.random.randint(0,2,m).astype(bool)

只是为了显示这些数据的样子：

In [2]: a
Out[2]: 
array([[ 0.,0.,0.],[ 0.,0.]])

In [3]: b
Out[3]: 
array([[ 1.,1.,1.],[ 1.,1.]])

In [4]: boolvec
Out[4]: array([ True,True,False,False],dtype=bool)

在这种情况下,为此使用np.where是最有效的.但是,我们需要boolvec的形状可以广播到与a和b相同的形状.因此,我们可以通过使用np.newaxis或None切片使它成为列向量(它们是相同的)：

In [5]: boolvec[:,None]
Out[5]: 
array([[ True],[ True],[False],[False]],dtype=bool)

然后我们可以使用np.where创建最终结果：

In [6]: c = np.where(boolvec[:,None],a,b)

In [7]: c
Out[7]: 
array([[ 0.,1.]])