[Numpy-discussion] simple indexing question

[josef.pktd at gmail.com]

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

>
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion at scipy.org
> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>