需要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)