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

Question

我有一组2d索引。

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

现在，我有一个不同的4维数组，arr，我想只提取一个数组（它是一个数组，因为它是4维），在indices数组中有相应的索引。它等同于以下代码。

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的地方使用indices_arr。

Answer 2

尝试使用numpy数组的take函数。你的代码应该是这样的：

outputarray= np.take(arr,indices)