根据存储在另一个数组或列表中的索引拆分 numpy 多维数组

Question

我有一个形状 = (12,2,3,3) 的 numpy 多维数组

import numpy as np
arr = np.arange(12*2*3*3).reshape((12,2,3,3))

我需要根据索引存储在另一个列表中的第二维来选择那些元素

indices = [0,1,0,0,1,1,0,1,1,0,1,1]

在一个数组中，其余的在另一个数组中。 任何一种情况下的输出都应该是另一个形状为 (12,3,3) 的数组

arr2 = np.empty((arr.shape[0],*arr.shape[-2:]))

我可以使用 for 循环来做到这一点

for i, ii in enumerate(indices):
    arr2[i] = arr[i, indices[ii],...]

但是，我正在寻找单班轮。

当我尝试使用列表作为索引进行索引时

test = arr[:,indices,...]

我得到形状为 (12,12,3,3) 而不是 (12,3,3) 的 test。你能帮帮我吗？

Answer 1

您可以使用np.arange 来索引第一个维度：

test = arr[np.arange(arr.shape[0]),indices,...]

或者只是 python range 函数：

test = arr[range(arr.shape[0]),indices,...]