[Numpy-discussion] simple indexing question
Neal Becker
ndbecker2 at gmail.com
Wed Sep 23 13:07:57 EDT 2009
josef.pktd at gmail.com wrote:
> On Wed, Sep 23, 2009 at 11:12 AM, Neal Becker <ndbecker2 at gmail.com> wrote:
>> Robert Cimrman wrote:
>>
>>> Neal Becker wrote:
>>>> I have an array:
>>>> In [12]: a
>>>> Out[12]:
>>>> array([[0, 1, 2, 3, 4],
>>>> [5, 6, 7, 8, 9]])
>>>>
>>>> And a selection array:
>>>> In [13]: b
>>>> Out[13]: array([1, 1, 1, 1, 1])
>>>>
>>>> I want a 1-dimensional output, where the array b selects an element
>>>> from each column of a, where if b[i]=0 select element from 0th row of a
>>>> and if b[i]=k select element from kth row of a.
>>>>
>>>> Easy way to do this? (Not a[b], that gives 5x5 array output)
>>>
>>> It might be stupid, but it works...
>>>
>>> In [51]: a
>>> Out[51]:
>>> array([[0, 1, 2, 3, 4],
>>> [5, 6, 7, 8, 9]])
>>>
>>> In [52]: b = [0,1,0,1,0]
>>>
>>> In [53]: a.T.flat[a.shape[0]*np.arange(a.shape[1])+b]
>>> Out[53]: array([0, 6, 2, 8, 4])
>>>
>>> cheers,
>>> r.
>>
>> Thanks. Is there really no more elegant solution?
>
> How about this?
>
>>>> a
> array([[0, 1, 2, 3, 4],
> [5, 6, 7, 8, 9]])
>>>> b
> array([0, 1, 0, 1, 0])
>
>>>> a[b,np.arange(a.shape[1])]
> array([0, 6, 2, 8, 4])
>
> Josef
>
Thanks, that's not bad. I'm a little surprised that given the fancy
indexing capabilities of np there isn't a more direct way to do this. I'm
still trying to wrap my mind around the fancy indexing stuff.
More information about the NumPy-Discussion
mailing list