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

[Keith Goodman]

On Thu, Aug 5, 2010 at 10:12 AM, Martin Spacek <numpy at mspacek.mm.st> wrote:
> 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])

Here's one way:

>> a.flat[i + a.shape[1] * np.arange(a.shape[0])]
   array([0, 3, 5, 6, 9])