Numpy 使用 integer 索引和 boolean 索引进行双片分配

Question

我已经知道 Numpy 带有精美索引的“双切片”创建副本而不是视图，解决方案似乎是将它们转换为单个切片（例如This question ）。 但是，我正面临这个特殊问题，我需要处理 integer 索引，然后是 boolean 索引，我不知所措。 问题（简化）如下：

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

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

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)

现在将掩码重塑为广播索引并将其用于 select

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

该作业现在有效，但我不确定这是否是您想要的确切作业。

Answer 2

好吧，显然我让事情变得复杂了。 无需组合索引。 下面的代码优雅地解决了这个问题：

b = a[..., idx_y, idx_x]
b[mask] = 1
a[..., idx_y, idx_x] = b
print(a[..., idx_y, idx_x][mask]) # all 1s

Answer 3

编辑：使用@Kevin 的解决方案，它实际上得到了正确的尺寸！

---

我没有在您的示例代码上专门尝试过，但我之前遇到过类似的问题。 我想我通过将掩码应用于索引来解决它，例如：

a[..., idx_y[mask], idx_x[mask]] = 1

- 这样，numpy 可以正确地将值分配给a数组。

---

EDIT2：发布一些测试代码作为评论删除格式。

a = np.arange(27).reshape([3, 3, 3])
ind_x = np.array([[0, 0], [1, 2]])
ind_y = np.array([[1, 2], [1, 1]])
x = np.broadcast_to(ind_x.ravel(), (4, 4))
y = np.broadcast_to(ind_y.ravel(), (4, 4)).T
# x1, y2 = np.meshgrid(ind_x, ind_y)  # above should be the same as this
mask = a[:, ind_y, ind_x] % 2 == 0  # what should this reshape to?
# a[..., x[mask], y[mask]] = 1  # Then you can mask away (may also need to reshape a or the masked x or y)