[FIXED] How to quickly grab specific indices from a numpy array?

Issue

But I don't have the index values, I just have ones in those same indices in a different array. For example, I have

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

Is there some NumPy method than can quickly look at both of these and extract all values from a whose indices match the indices of all 1's in b? I want it to result in:

array([4,6])

It is probably worth mentioning that my a array is multidimensional, while my b array will always have values of either 0 or 1. I tried using NumPy's logical_and function, though this returns ValueError with a and b having different dimensions:

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,) 

Though this method does seem to work if a is flat. Either way, the return type of numpy.logical_and() is a boolean, which I do not want. Is there another way? Again, in the second example above, the desired return would be

array([[4,5]])

Obviously I could write a simple loop to accomplish this, I'm just looking for something a bit more concise.

Edit:

This will introduce more constraints, I should also mention that each element of the multidimensional array a may be any arbitrary length, that does not match its neighbour.

Solution

You can simply use fancy indexing.

b == 1

will give you a boolean array:

>>> 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)

which you can pass as an index to a.

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

Demo for your second example:

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

Answered By - timgeb