Numpy: вещание + логическое индексирование

Question

У меня есть следующие массивы NumPy

A: shape (n1, n2) array of float
B: shape (n2,) array of float
M: shape (n1, n2) array of bool

Как превратить следующий псевдо-код в эффективный реальный код? Массивы могут быть огромными, возможно,> 100 миллионов элементов.

A[M] = ("B broadcast to shape (n1,n2)")[M]

Answer 1

Вещание просто и эффективно использует память:

A, B, M = np.broadcast_arrays(A, B, M)

Однако, используя это B в вашем коде A[M] = B[M] не будет эффективным использованием памяти, потому что B[M] имеет столько реальных элементов, сколько M имеет True ценности.

Вместо этого используйте:

np.putmask(A, M, B)

поскольку B повторяется автоматически с putmask функция, вы не должны даже транслировать ее. Хотя, думаю, это не помешает.