Indexing NumPy 2D array with another 2D array

Question

I have something like

m = array([[1, 2],
            [4, 5],
            [7, 8],
            [6, 2]])

and

select = array([0,1,0,0])

My target is

result = array([1, 5, 7, 6])

I tried _ix as I read at Simplfy row AND column extraction, numpy, but this did not result in what I wanted.

p.s. Please change the title of this question if you can think of a more precise one.

Answer 1

The numpy way to do this is by using np.choose or fancy indexing/take (see below):

m = array([[1, 2],
           [4, 5],
           [7, 8],
           [6, 2]])
select = array([0,1,0,0])

result = np.choose(select, m.T)

So there is no need for python loops, or anything, with all the speed advantages numpy gives you. m.T is just needed because choose is really more a choise between the two arrays np.choose(select, (m[:,0], m[:1])), but its straight forward to use it like this.

---

Using fancy indexing:

result = m[np.arange(len(select)), select]

And if speed is very important np.take, which works on a 1D view (its quite a bit faster for some reason, but maybe not for these tiny arrays):

result = m.take(select+np.arange(0, len(select) * m.shape[1], m.shape[1]))