numpy: broadcasting ndarray of booleans

Question

У меня есть ситуация, когда у меня есть ndarray X floats, скажем, 100x10, и я хочу посмотреть некоторые условия на первый столбец и создать логический ndarray B формы 100x1. Затем я хочу использовать B как индекс в X, чтобы вытащить значения, где находится True. Однако для каждого True в B я хочу вытащить всю строку X. Я думал, что это будет работать автоматически, так как B будет транслироваться до формы 100x10. Однако похоже, что это не так. Вот пример использования 2x2 и 2x1 ndarrays.

a = np.array([True, False]) 
a.shape = (2,1) 
b = np.array([1, 2, 3, 4]) 
b.shape = (2,2) 
print(a) 
print(b) 
print(b[a]) 

Печатается

[[True] 
[False]] 

[[ 1 2 ] 
[ 3 4 ]] 

[1] 

Я ожидал, что это напечатать [1 2]. Почему трансляция работает так, как я ожидаю?

Answer 1

Правила так называемого «индексирования индексирования» - detailed here. В частности, когда индекс, obj, является массив NumPy из DTYPE bool, x[obj]

... всегда эквивалентна (но быстрее, чем) х [obj.nonzero()] , где, как описано выше , obj.nonzero() возвращает кортеж (длины obj.ndim) массивов целых индексов, показывающих True элементов obj.

Поскольку

In [4]: a.nonzero() 
Out[4]: (array([0]), array([0])) 

b[a] эквивалентно b[a.nonzero()], который

In [6]: b[(np.array([0]), np.array([0]))] 
Out[6]: array([1]) 
In [7]: b[a] 
Out[7]: array([1]) 

---

Если вы хотите использовать булево массив a выбрать строки из b, то , как утверждает Йоран Бисли, jus т держать a как 1-мерного булева массива:

import numpy as np 

a = np.array([True, False]) 
b = np.array([1, 2, 3, 4]) 
b.shape = (2,2) 
print(b[a])  
# [[1 2]]