我有两个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.]])