如何从numpy数组中快速获取特定索引？

Question

但我没有索引值，我只是在一个不同的数组中有相同的索引值。例如，我有

a = array([3,4,5,6]) 
b = array([0,1,0,1]) 

有一些NumPy相比能快速了解一下这两个从a其索引匹配所有1的在b索引中提取所有值？我想这导致：

array([4,6]) 

这可能是值得一提的是我的a数组是多维的，而我的b阵列将始终有两种0或1的值。我试着用NumPy的logical_and功能，虽然这将返回ValueError与a和b具有不同尺寸：

a = numpy.array([[3,2], [4,5], [6,1]]) 
b = numpy.array([0, 1, 0]) 
print numpy.logical_and(a,b) 

ValueError: operands could not be broadcast together with shapes (3,2) (3,) 

虽然这种方法似乎工作，如果a是平的。无论哪种方式，返回类型numpy.logical_and()是一个布尔值，我不想要。有另一种方法吗？同样，在上面的第二个例子中，所需的回报将是

array([[4,5]]) 

很显然，我可以写一个简单的循环来做到这一点，我只是在寻找一些更简洁。

编辑：

这将引入更多的限制，我还应该提到的是，多维数组a中的每个元素可以是任意长度，即不匹配其邻居。

Answer 1

您可以简单地使用花式索引。

b == 1 

会给你一个布尔数组：

>>> from numpy import array 
>>> a = array([3,4,5,6]) 
>>> b = array([0,1,0,1]) 
>>> b==1 
array([False, True, False, True], dtype=bool) 

其中可以传递作为索引到。

>>> a[b==1] 
array([4, 6]) 

演示你的第二个例子：

>>> a = array([[3,2], [4,5], [6,1]]) 
>>> b = array([0, 1, 0]) 
>>> a[b==1] 
array([[4, 5]]) 

Answer 2

你可以使用compress：

>>> a = np.array([3,4,5,6]) 
>>> b = np.array([0,1,0,1]) 
>>> a.compress(b) 
array([4, 6]) 

您可以为多维情况下的axis说法：

>>> a2 = np.array([[3,2], [4,5], [6,1]]) 
>>> b2 = np.array([0, 1, 0]) 
>>> a2.compress(b2, axis=0) 
array([[4, 5]]) 

即使您正在索引的a的轴与b的长度不同，该方法也可以工作。