python - Numpy 中的多维索引

Question

我检查了 Numpy docs所以。对此，但我找不到答案。也许做不到。

基本上我有一个形状为 (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]])