Numpy 中的多维索引

Question

我检查了Numpy docs 和 s.o.对此，但我找不到答案。也许做不到。

基本上我有一个形状为 (3,4,5) 的数组 probabilities。在第 3 维中有 5 个元素，它们的总和为 1。我想要的第 3 维中的元素索引对应于形状 (3,4) 的数组index 中的值。有意义吗？

所以如果probabilities[0,0,:] 等于[0.1, 0.1, 0.2, 0.4, 0.2] 并且index[0,0] 等于2，那么我想要第三个元素0.2。

我尝试了probabilities[index] 和其他东西，但没有运气。

这可以不用循环来完成吗？

Answer 1

制作一个样本数组：

In [291]: A = np.arange(2*3*4).reshape(2,3,4)
In [292]: A[0,0,:]
Out[292]: array([0, 1, 2, 3])
In [293]: A[0,0,2]
Out[293]: 2

制作一个示例 idx：

In [294]: idx = np.random.randint(0,4,(2,3),int)
In [295]: idx
Out[295]: 
array([[0, 3, 0],
       [1, 0, 1]])

这些是第 3 维的索引值。制作数组以在第 2 个维度上进行索引：

In [299]: I,J=np.ix_(np.arange(A.shape[0]),np.arange(A.shape[1]))
In [300]: I,J
Out[300]: 
(array([[0],
        [1]]), array([[0, 1, 2]]))
In [301]: A[I,J,idx]
Out[301]: 
array([[ 0,  7,  8],
       [13, 16, 21]])

测试：

In [302]: A[0,1,3]
Out[302]: 7
In [304]: A[1,2,1]
Out[304]: 21

有多种方法可以获取这些I,J。 np.ix_ 很简单。 np.ogrid、np.mgrid 甚至 np.meshgrid 也是如此。

In [306]: I,J = np.mgrid[0:2,0:3]
In [307]: I,J
Out[307]: 
(array([[0, 0, 0],
        [1, 1, 1]]), array([[0, 1, 2],
        [0, 1, 2]]))
In [308]: A[I,J,idx]
Out[308]: 
array([[ 0,  7,  8],
       [13, 16, 21]])