Numpy中的多維索引

Question

我檢查了Numpy文檔 ，依此類推，但找不到答案。 也許無法完成。

基本上我有一個形狀為（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]])

這些是第三維的索引值。 制作要在第1個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也是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]])