python - 如何用掩码和索引替换双 for 循环？

Question

我有几个嵌套的 for 循环可以做正确的事情(数组的屏蔽副本)。然而性能太慢了，我觉得必须有更好的 Pythonic 方法来做到这一点。目标是使用掩码来确定何时从源复制数据，使用坐标作为源的索引。有效的循环代码如下:

import numpy as np
dest = np.zeros((4,4,2))
source = range(32)
source = np.reshape(source,(4,4,2))
mask = np.ones((4,4),bool)
mask[1,0] = 0
coord = np.ones((4,4,2),int)

for y in range (0,dest.shape[0]):
    for x in range (0, dest.shape[1]):
        if np.all(mask[y,x]):
            dest[y,x] = source[coord[y,x,0], coord[y,x,1]]

print dest

运行后，dest看起来像这样:

[[[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[  0.   0.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]
 [[ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]
  [ 10.  11.]]]

source[1,1] 被复制到所有 dest，但 dest[1,0] 除外，因为 mask [1,0] 设置为 False。 mask 的其余部分为 True。谁能告诉我如何用更有效的东西替换循环？

Answer 1

使用numpy.where 。您必须向 mask 添加一个额外的维度，这样它就会 broadcast .

dest = np.where(mask[:,:,None], source[coord[:,:,0], coord[:,:,1]], dest)

或者如果合适的话:

dest = np.where(mask[:,:,None], source[coord[:,:,0], coord[:,:,1]], np.zeros((4,4,2)))