python - 多个索引的 Numpy 数组替换为不同的矩阵

Question

我有一个二维索引数组。

indices = [[2,4], [6,77], [102,554]]

现在，我有一个不同的 4 维数组 arr，我只想提取一个数组(它是一个数组，因为它是 4 维的)在索引数组中具有相应的索引。相当于下面的代码。

for i in range(len(indices)):
    output[i] = arr[indices[i][0], indices[i][1]]

但是，我意识到使用显式 for 循环会产生缓慢的结果。有没有我可以使用的内置 numpy API？此时，我尝试使用 np.choose、np.put、np.take，但没有成功产生我想要的结果。谢谢!

Answer 1

我们需要用索引中的两列索引到前两个轴(将其视为数组)。

因此，只需转换为数组和索引，就像这样 -

indices_arr = np.array(indices)
out = arr[indices_arr[:,0], indices_arr[:,1]]

或者我们可以直接提取那些而不转换为数组然后索引 -

d0,d1 = [i[0] for i in indices], [i[1] for i in indices]
out = arr[d0,d1]

另一种提取元素的方法是转换为元组，就像这样 -

out = arr[tuple(indices_arr.T)]

如果 indices 已经是一个数组，跳过转换过程并在我们有 indices_arr 的地方使用 indices。