[Numpy-discussion] use index array of len n to select columns of n x m array

Martin Spacek numpy@mspacek.mm...
Thu Aug 5 12:12:50 CDT 2010


I want to take an n x m array "a" and index into it using an integer index array 
"i" of length n that will pull out the value at the designated column from each 
corresponding row of "a".

>>> a = np.arange(10)
>>> a.shape = 5, 2
>>> a
array([[0, 1],
        [2, 3],
        [4, 5],
        [6, 7],
        [8, 9]])
>>> i = np.array([0, 1, 1, 0, 1])

I want:

>>> b = a.foo(i)
>>> b
array([0, 3, 5, 6, 9])

What's foo? I can't get take() to do what I want. I went and wrote my own little 
Cython function to do this, but that seems silly (and is also array dtype 
dependent). I've tried reading through the numpy book, and I'm sure this is 
somewhere on the list, but I can't find it. I think it has something to do with 
fancy indexing. I should know how to do this by know...

Cheers,

Martin



More information about the NumPy-Discussion mailing list