python - Numpy 双切片赋值，整数索引后跟 bool 索引

Question

我已经知 Prop 有奇特索引的 Numpy“双切片”会创建副本而不是 View ，解决方案似乎是将它们转换为单个切片(例如 This question )。然而，我面临着这个特殊的问题，我需要处理整数索引，然后处理 bool 索引，但我不知道该怎么办。问题(简化)如下:

a = np.random.randn(2, 3, 4, 4)
idx_x = np.array([[1, 2], [1, 2], [1, 2]])
idx_y = np.array([[0, 0], [1, 1], [2, 2]])
print(a[..., idx_y, idx_x].shape) # (2, 3, 3, 2)
mask = (np.random.randn(2, 3, 3, 2) > 0)
a[..., idx_y, idx_x][mask] = 1 # assignment doesn't work

如何才能完成作业？

Answer 1

不确定，但一个想法是手动进行广播并分别添加掩码，就像蒂姆建议的那样。 idx_x 和 idx_y 都具有相同的形状 (3,2)，它们将被广播到形状 (6,6) 来自笛卡尔积 (3*2)^2。

x = np.broadcast_to(idx_x.ravel(), (6,6))
y = np.broadcast_to(idx_y.ravel(), (6,6))

# this should be the same as
x,y = np.meshgrid(idx_x, idx_y)

现在将掩码 reshape 为广播索引并使用它来选择

mask = mask.reshape(6,6)
a[..., x[mask], y[mask]] = 1

该作业现在可以运行，但我不确定这是否是您想要的作业。