如何從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的長度不同，該方法也可以工作。