Indexing NumPy 2D array with another 2D array

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

I prefer to use NP.where for indexing tasks of this sort (rather than NP.ix_)

What is not mentioned in the OP is whether the result is selected by location (row/col in the source array) or by some condition (e.g., m >= 5). In any event, the code snippet below covers both scenarios.

Three steps:

  1. create the condition array;

  2. generate an index array by calling NP.where, passing in this condition array; and

  3. apply this index array against the source array


>>> import numpy as NP

>>> cnd = (m==1) | (m==5) | (m==7) | (m==6)
>>> cnd
  matrix([[ True, False],
          [False,  True],
          [ True, False],
          [ True, False]], dtype=bool)

>>> # generate the index array/matrix 
>>> # by calling NP.where, passing in the condition (cnd)
>>> ndx = NP.where(cnd)
>>> ndx
  (matrix([[0, 1, 2, 3]]), matrix([[0, 1, 0, 0]]))

>>> # now apply it against the source array   
>>> m[ndx]
  matrix([[1, 5, 7, 6]])


The argument passed to NP.where, cnd, is a boolean array, which in this case, is the result from a single expression comprised of compound conditional expressions (first line above)

If constructing such a value filter doesn't apply to your particular use case, that's fine, you just need to generate the actual boolean matrix (the value of cnd) some other way (or create it directly).