Numpy: индексация массивов

Question

Предоставлено следующий пример

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

Чтобы получить суб-массив, содержащий 1 в первом столбце:

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

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

Как получить (без петель) суб-массив содержащие 1 и 5? Должно быть не что-то вроде

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

который не работает?

Answer 1

Способ № 1: использование побитового или | совместить условие:

>>> 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]]) 

Способ № 2: использовать np.in1d, который, вероятно, легче, если есть много значений:

>>> 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]])