CODEWITHSUNDEEP
arraysnumpy
Mozzie
Mozzie

Reputation: 353

Select columns in numpy

Assumes there is an index and a matrix L

>>> index
(array([0, 2, 3, 3]), array([0, 2, 2, 3]))
>>> L
array([[  1,  -1,  -5, -10],
   [-15,   0,  -1,  -5],
   [-10, -15,  10,  -1],
   [ -5, -10,   1,  15]])

I want to select the columns according to the index[1], I've tried:

>>> L[:,index[1]]
array([[  1,  -5,  -5, -10],
   [-15,  -1,  -1,  -5],
   [-10,  10,  10,  -1],
   [ -5,   1,   1,  15]])

but the result is not i expected, what I expected is:

>>> for i in index[1]:
...     print L[:,i]
[  1 -15 -10  -5]
[-5 -1 10  1]
[-5 -1 10  1]
[-10  -5  -1  15]

How can i get the expected result without for loop? and why this unexpected result comes out? Thanks.

Upvotes: 2

Views: 538

Answers (1)

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 477684

You simply need to transpose it:

L[:,index[1]].T
#             ^ transpose

By using a transpose, the columns are rows and vice versa. So here (you can transpose before the selection, and then use L.T[index[1],:]) we first make the selection and then turn the columns into rows.

This produces:

>>> L[:,index[1]].T
array([[  1, -15, -10,  -5],
       [ -5,  -1,  10,   1],
       [ -5,  -1,  10,   1],
       [-10,  -5,  -1,  15]])

Note that of course behind the curtains there are still some loops that are done. But these are done outside Python and thus are more efficient.

Upvotes: 3

Related Questions