元组多重索引

Question

我有一个多维度的

np.array

。我知道前 N 维和后 M 维的形状。例如，

>>> n = (3,4,5)
>>> m = (6,)
>>> a = np.ones(n + m)
>>> a.shape
(3, 4, 5, 6)

使用元组作为索引可以像前 N 维一样快速索引，比如

>>> i = (1,1,2)
>>> a[i].shape
(6,)

使用列表并不能得到我需要的相同结果

>>> i = [1,1,2]
>>> a[i].shape
(3, 4, 5, 6)

但是我在进行多重索引（检索/分配值）时遇到麻烦。例如，

>>> i = (1,1,2)
>>> j = (2,2,2)

我需要传递类似的东西

>>> a[[i, j]]

并获得输出形状

(2, 6)

。

相反我得到了

>>> a[[i, j]].shape
(2, 3, 4, 5, 6)

或

>>> a[(i, j)].shape
(3, 5, 6)

我总是可以循环或更改索引的方式（例如使用

np.reshape

和

np.unravel_index

），但是有没有更 Pythonic 的方法来实现我所需要的？

Answer 1

提取每个选择，然后将它们重新拼接到一个新数组中？

>>> np.array([a[i], a[j]]).shape
(2, 6)