如何从两个数组中选择项？

Question

我有两个数组a和b，其中包含一些数据。我还有一个数组select，我想用它从a或b中进行选择。我只是想知道是否有一个pythonic方法可以这样做。下面是我当前的实现，它将a和b的每一行放入一个列表中，然后从中进行选择。你知道吗

a = np.zeros(shape=(10,2,1,3,4))
b = np.ones(shape=(10,2,1,3,4))

select = [1,1,1,0,1,0,1,0,1,0]

c = []
for a1,b1,select1 in zip(a,b,select):
    a1b1 = [a1,b1]
    c.append(a1b1[select])

Answer 1

如果仅从两个数组中选择，则可以使用select作为权重：

w = np.array(select)
c = (a.T * w + b.T * (1 - w)).T

如果要组合两个以上的数组，则此方法不起作用。你知道吗

Answer 2

np.where是一个非常有效的解决方案，但最多只能用于两个数组。你知道吗

我会这样做：

np.stack((a, b))[select, np.arange(len(a))]

这应该扩展到任意数量的阵列，例如：

a = ...
b = ...
# ...
z = ...
select = [0, 13, 2, 5, 25, ...]
np.stack((a, b, ... z))[select, np.arange(len(a))]

Answer 3

In [512]: a = np.zeros(shape=(10,2,1,3,4)) 
     ...: b = np.ones(shape=(10,2,1,3,4)) 
     ...:  
     ...: select = [1,1,1,0,1,0,1,0,1,0] 

以@DYZ相同的方式使用转置：

In [513]: res = np.where(select, a.T, b.T).T                                    
In [514]: res.shape                                                             
Out[514]: (10, 2, 1, 3, 4)

如果不想转置，可以向select添加维度，使其广播到与a和b相同的形状：

In [516]: res1 = np.where(np.array(select)[:,None,None,None,None],a,b)          
In [517]: res1.shape                                                            
Out[517]: (10, 2, 1, 3, 4)
In [518]: np.allclose(res,res1)                                                 
Out[518]: True