使用单个索引项的numpy子矩阵3D

Question

我有一个3D numpy数组：

K = (np.arange(36)).reshape((4,3,3))+1
[[[ 1  2  3]
[ 4  5  6]
[ 7  8  9]]

[[10 11 12]
[13 14 15]
[16 17 18]]

[[19 20 21]
[22 23 24]
[25 26 27]]

[[28 29 30]
[31 32 33]
[34 35 36]]]

其中K中的每个项目都是矩阵。 现在，我想使用某个索引向量

获取所有2D子矩阵

我知道这有可能：

idx = np.s_[:,:2,:2]
K_sub = K[idx]
[[[ 1  2]
[ 4  5]]

[[10 11]
[13 14]]

[[19 20]
[22 23]]

[[28 29]
[31 32]]]

问题是我想使用任意索引数组而不是切片来选择行和列。

此外，我想使用单个对象来获取子矩阵列表，例如：

K_sub = [magic_indexing]

而不是：

K_sub = np.array([k_[train][:,train] for k_ in K])

存在一种简单的方法吗？

Answer 1

不确定它是否足够简单，但有一种方法是使用np.ix_，因此使用advanced-indexing，就像这样 -

K[np.ix_(np.arange(K.shape[0]), train, train)]