[Numpy-discussion] simple indexing question

Neal Becker ndbecker2 at gmail.com
Wed Sep 23 11:12:51 EDT 2009


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?




More information about the NumPy-Discussion mailing list