Numpy：广播+布尔索引

Question

我有以下 numpy 数组

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

如何将以下伪代码转换为高效真实代码？数组可能很大，可能有超过 1 亿个元素。

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

Answer 1

广播简单且节省内存：

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

但是，在您的代码 A[M] = B[M] 中使用此 B 不会节省内存，因为 B[M] 的实际元素数量与 M 的 True 值一样多。

改为使用：

np.putmask(A, M, B)

由于B 使用putmask 函数自动重复，您甚至不必广播它。虽然我想这样做不会有什么坏处。