3dアレイ上で効率的なナンシースライス

Question

私は3D numpy配列のスライスを行う最も効率的な方法を探しています。これはデータのサブセットで、テスト目的のみのものです。

in_arr =np.array([[[0,1,2,5],[2,3,2,6],[0,1,3,2]],[[1,2,3,4],[3,1,0,5],[2,4,0,1]]]) 
indx =[[3,1,2],[2,0,1]] 

indxで値を取得する必要があります。たとえば、indx [0] [0]は3なので、in_arr [0] [0]の3番目の要素、この場合は5を探しています。

私は次のコードを持っています。それが必要ですが、時間複雑度はn^2です。これは私が喜んでいません。

list_in =[] 
for x in range(len(indx)): 
    arr2 = [] 
    for y in range(len(indx[x])): 
     arr2.append(in_arr[x][y][indx[x][y]]) 
     #print in_arr[x][y][indx[x][y]] 
    list_in.append(arr2) 
print list_in 

大規模なデータセットに対して同じタスクを実行するための非常に高速で効率的な方法を探しています。

Answer 1

ブロードキャストされたインデックスの配列を使用すると、これを効率的に行うことができます。例えば：ここに何numpyの

i1 = np.arange(2)[:, np.newaxis] 
i2 = np.arange(3)[np.newaxis, :] 
i3 = np.array(indx) 
in_arr[i1, i2, i3] 
# array([[5, 3, 3], 
#  [3, 3, 4]]) 

は、効果的に3つの索引アレイのエントリと一致し、in_arrから関連する項目を抽出することである：[:, np.newaxis]と[np.newaxis, :]用語の理由は、それが3つの配列を再形成することですnumpyのbroadcastingルールで互換性があります。