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]]) 

これらは第三次元のインデックス値です。第一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]])