Numpy 使用二维整数数组索引三维数组的最后一个维度

我有一个3D

data=NxMxD

numpy数组和另一个2D

idx=NxM

整数数组，其值在

[0，D-1]

范围内。我想在该位置的

idx

数组给定的深度处对每个

data=NxM

条目执行基本更新

例如，对于

N=M=D=2

：

data = np.zeros((2,2,2))
idx = np.array([[0,0],[1, 1]], int)

我想执行一个简单的操作，比如：

data[..., idx] += 1

我的预期产出是：

>>> data
array([[[ 1.,  0.],
        [ 1.,  0.]],

       [[ 0.,  1.],
        [ 0.,  1.]]])

idx

为每个二维坐标指示应更新的

D

。上述操作不起作用

我在中找到了通过使用以下方法解决索引问题的SO：

data[np.arange(N)[:, None], np.arange(M)[None, :], idx] += 1

它工作得很好，但是看起来非常可怕，需要手动索引整个矩阵，这似乎是一个非常简单的操作（使用一个矩阵作为最后一个通道的索引掩码）

---

有更好的解决方案吗？

使用

numpy.ix

看起来并不可怕，但使用花式索引的基本思想仍然是一样的

x = np.arange(N)
y = np.arange(M)

xx,yy = np.ix_(x,y)

data[xx,yy,idx] += 1

注意

问题是您想更改

数据的值。如果您只想获得符合
idx
的值，您可以这样做

out = np.choose(idx,data.transform(2,0,1))

但是，这会为您提供
数据
值的副本，而不是视图，这意味着

out += 1

对
数据中的值没有影响