[Numpy-discussion] simple indexing question

Robert Cimrman cimrman3 at ntc.zcu.cz
Wed Sep 23 09:59:24 EDT 2009


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.




More information about the NumPy-Discussion mailing list