numpy - 使用 2D 整数数组索引 3D 数组的最后一个维度

Question

我有一个 3D data = NxMxD numpy 数组，和另一个 2D idx = NxM 整数数组，其值在 [0, D- 范围内1]。我想对每个 data = NxM 条目执行基本更新，更新深度由该位置的 idx 数组指定。

例如，对于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 指示每个 2D 坐标应更新哪个 D。上面的操作不行。

我找到了this approach在 SO 中，它通过使用以下方法解决了索引问题:

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

它工作正常，但看起来非常可怕，需要手动索引整个矩阵，这似乎是一个非常简单的操作(使用一个矩阵作为最后一个 channel 的索引掩码)。

有没有更好的解决办法？

Answer 1

使用numpy.ix_，它看起来并不那么可怕，但使用花式索引的基本思想仍然是相同的

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

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

data[xx,yy,idx] += 1

注意

问题是您想要更改data的值。如果您只想根据 idx 获得值，您可以这样做

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

但是，这为您提供了 data 值的副本，而不是 View ，这意味着

out += 1

对数据中的值没有影响。