需要 Numpy 花式索引建议

Question

首先，很抱歉标题描述不够完美。

我有一个 (500, 500, 500) ndarray "A" 的浮点数，还有另一个 (500, 500, 500) ndarray "Ind" 的整数。

我想形成第三个 (500, 500, 500) ndarray "Anext"，其中 "Anext" 中索引为 (x, y, z) 的条目是 A[B[x, y, z], y , z]。

我已经研究过索引，但我发现很难找到一个有用的解决方案。感谢您的帮助。

Answer 1

使用numpy.indices()创建轴1和轴2的索引数组，这里是一个例子：

import numpy as np

a = np.random.rand(4, 4, 4)
ind = np.random.randint(0, 4, size=(4, 4, 4))

i0, i1, i2 = np.indices(a.shape)

b = a[ind, i1, i2]

检查结果：

for (i, j, k), v in np.ndenumerate(b):
    assert v == a[ind[i, j, k], j, k]

i0、i1、i2的形状为(4, 4, 4)，为了节省一些内存，也可以使用numpy.ix_()：

i0, i1, i2 = np.ix_(*[np.arange(d) for d in a.shape])

其中i1.shape、i2.shape 是(1, 4, 1), (1, 1, 4)