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

时间:2015-05-28 16:34:32

标签: numpy indexing mask

我有一个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。上述操作不起作用。

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

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

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

有没有更好的解决方案?

1 个答案:

答案 0 :(得分: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值的副本,而不是视图,这意味着

out += 1

data中的值没有影响。