使用3D数组的索引填充4D数组

Question

我有两个3D数组，一个包含我正在使用的值，另一个包含索引。我想使用这两个填充4D数组。

索引数组的每个条目都指向输入数组的一行。4

起初，我只是简单地遍历i，j和k的值，然后手动填充每一行。但是，由于这是一个机器学习项目，因此此方法花费的时间太长。

# x.shape = (8, 2500, 3)
# ind.shape = (8, 2500, 9)

M = np.empty(8, 2500, 9, 3)

for i in range(0, M.shape[0]):
    for j in range(0, M.shape[1]):
        for k in range(0, M.shape[2]):
            M[i, j, k, :] = x[i, ind[i, j, k], :]

是否存在一种更快的方法？

Answer 1

您可以尝试以下操作：

import numpy as np
M = x[np.arange(0,ind.shape[0])[:, None, None], ind]

需要[:, None, None]才能将np.arange(0,ind.shape[0])广播到正确的维度以索引数组x。

作为测试，您可以使用当前方法生成数组M，然后使用上述方法生成数组M_，并确认(M == M_).all()返回{{1} }。

我将其速度至少提高了30倍。