up vote 4 down vote favorite
2

I have a 2D numpy array (i.e matrix) A which contains useful data interspread with garbage in the form of column vectors as well as a 'selection' array B which contains '1' for those columns that are important and 0 for those that are not. Is there a way to select only those columns from A that correspond to ones in B? i.e i have a matrix

A = array([[ 0,  1,  2,  3,  4],   and a vector B = array([ 0,  1,  0,  1,  0])
           [ 5,  6,  7,  8,  9],
           [10, 11, 12, 13, 14],
           [15, 16, 17, 18, 19],
           [20, 21, 22, 23, 24]])

and I want

array([[1,   3],
       [6,   8],
       [11, 13],
       [16, 18],
       [21, 23]])

Is there an elegant way to do so? Right now i just have a for loop that iterates through B.

NOTE: the matrices that i'm dealing with are large, so i don't want to use numpy masked arrays, as i simply don't want the masked data

share|improve this question

3 Answers

up vote 6 down vote accepted 
>>> A
  array([[ 0,  1,  2,  3,  4],
         [ 5,  6,  7,  8,  9],
         [10, 11, 12, 13, 14],
         [15, 16, 17, 18, 19],
         [20, 21, 22, 23, 24]])
>>> B = NP.array([ 0,  1,  0,  1,  0])

>>> # convert the indexing array to a boolean array
>>> B = NP.array(B, dtype=bool)

>>> # index A against B--indexing array is placed after the ',' because
>>> # you are selecting columns

>>> res = A[:,B]

>>> res
  array([[ 1,  3],
         [ 6,  8],
         [11, 13],
         [16, 18],
         [21, 23]])


The syntax for index-based slicing in NumPy is elegant and simple. A couple of rules cover a majority of use cases:

  • the form is [rows, columns]

  • specify all rows or all columns using a colon ":" e.g., [:, 4] (extracts the entire 5th column)

share|improve this answer
this is great,thanks! – pratikm Dec 10 '11 at 1:18
up vote 2 down vote

Not sure if it's the most efficient way (because of the transposition), but it should be better than a for loop:

A.T[B == 1].T
share|improve this answer
A[:, B==1] works here. – Steve Tjoa Dec 10 '11 at 4:54
up vote 0 down vote

I was interested to do the same but to slice row & column using the boolean values of vector B, the solution was simple:

res = A[:,B][B,:]
share|improve this answer

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.