Numpy: Indexing of arrays

Question

Given the following example

d = array([[1, 2, 3],
           [1, 2, 3],
           [1, 3, 3],
           [4, 4, 4],
           [5, 5, 5]
          ])

To get the sub-array containing 1 in the first column:

d[ d[:,0] == 1 ]

array([[1, 2, 3],
       [1, 2, 3],
       [1, 3, 3]])

How to get (without loops) the sub-array containing 1 and 5? Shouldn't be something like

d[ d[:,0] == [1,5] ]  #   ---> array([1, 2, 3])

which does not work?

DSM · Accepted Answer · 2013-10-10 11:44:32Z

7

Method #1: use bitwise or | to combine the conditions:

>>> d
array([[1, 2, 3],
       [1, 2, 3],
       [1, 3, 3],
       [4, 4, 4],
       [5, 5, 5]])
>>> (d[:,0] == 1) | (d[:,0] == 5)
array([ True,  True,  True, False,  True], dtype=bool)
>>> d[(d[:,0] == 1) | (d[:,0] == 5)]
array([[1, 2, 3],
       [1, 2, 3],
       [1, 3, 3],
       [5, 5, 5]])

Method #2: use np.in1d, which is probably easier if there are a lot of values:

>>> np.in1d(d[:,0], [1, 5])
array([ True,  True,  True, False,  True], dtype=bool)
>>> d[np.in1d(d[:,0], [1, 5])]
array([[1, 2, 3],
       [1, 2, 3],
       [1, 3, 3],
       [5, 5, 5]])

answered Oct 10, 2013 at 11:44

DSM

355k67 gold badges606 silver badges504 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Tengis Over a year ago

the second solution is what I need. I was not aware of this in1d(). Thanks!

Collectives™ on Stack Overflow

Numpy: Indexing of arrays

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related