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

josef.pktd@gmai... josef.pktd@gmai...
Thu Aug 5 12:26:06 CDT 2010


On Thu, Aug 5, 2010 at 1:12 PM, Martin Spacek <numpy@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])
>
> 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...

like this ?

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

Josef

>
> Cheers,
>
> Martin
>
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion@scipy.org
> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>


More information about the NumPy-Discussion mailing list