Numpy 双切片赋值，整数索引后跟布尔索引

Question

我已经知道带有精美索引的 Numpy“双切片”会创建副本而不是视图，解决方案似乎是将它们转换为一个切片（例如）。但是，我遇到了这个特殊问题，我需要处理一个整数索引，然后是布尔索引，我不知道该怎么做。问题（简化）如下：

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

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

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

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

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

EDIT2：Post 一些测试代码作为注释删除格式。

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)

Answer 2

不确定，但一个想法是手动进行广播并分别添加 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 3

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

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

Numpy 双切片赋值，整数索引后跟布尔索引

Numpy double-slice assignment with integer indexing followed by boolean indexing

python

numpy

slice