Numpy for文を使わず行列内にあるインデックス値を元に別の行列の中から要素を取り出したい

Question

前提

機械学習の前処理としてNumpyで行列操作をしようとしています
3次元行列idxにインデックス値、３次元配列aに取り出したい数値が入っています

idxの形状は
データ数×１つのデータの中から取り出したい要素数×2(取り出したい要素がある座標のXとYの値)

aの形状は
データ数×Y方向の大きさ×X方向の大きさ

実現したいこと

現在以下のようにfor文を使って処理をしているのですが非常に時間がかかります
質問の為に簡略化をして以下のコードではループの数を3にしていますが、実際は数十万程度の回数ループを行います
これをFor文を使わない等の方法で高速化をする方法はありますでしょうか

ソースコード

import numpy as np


idx = np.array([[[0, 0], [0, 1], [0, 2], [0, 3]],
                [[1, 0], [1, 1], [1, 2], [1, 3]],
                [[2, 0], [2, 1], [2, 2], [2, 3]]])

a = np.arange(90).reshape(3,3,10)
b = np.empty([3, 4])
for i in range(3):
    b[i,:] = a[i, idx[i, :, 0], idx[i, :, 1]]

実行結果

idx
[[[0 0]
  [0 1]
  [0 2]
  [0 3]]

 [[1 0]
  [1 1]
  [1 2]
  [1 3]]

 [[2 0]
  [2 1]
  [2 2]
  [2 3]]]
-----------------------------------------
a
[[[ 0  1  2  3  4  5  6  7  8  9]
  [10 11 12 13 14 15 16 17 18 19]
  [20 21 22 23 24 25 26 27 28 29]]

 [[30 31 32 33 34 35 36 37 38 39]
  [40 41 42 43 44 45 46 47 48 49]
  [50 51 52 53 54 55 56 57 58 59]]

 [[60 61 62 63 64 65 66 67 68 69]
  [70 71 72 73 74 75 76 77 78 79]
  [80 81 82 83 84 85 86 87 88 89]]]
-----------------------------------------
b
[[ 0.  1.  2.  3.]
 [40. 41. 42. 43.]
 [80. 81. 82. 83.]]

Answer 1

i = np.arange(idx.shape[0]).reshape(-1, 1)
b = a[i, idx[:, :, 0], idx[:, :, 1]]

0次元目のインデックスはnumpy.arange で作れるのですが、他の次元とshapeを合わせる必要があるので、reshape します。
reshape で (3,) から (3, 1) に変換して、a[...] の処理のときにブロードキャストで (3, 4) になっています。

Answer 2

import numpy as np


idx = np.array([[[0, 0], [0, 1], [0, 2], [0, 3]],
                [[1, 0], [1, 1], [1, 2], [1, 3]],
                [[2, 0], [2, 1], [2, 2], [2, 3]]])

a = np.arange(90).reshape(3,3,10)
b = np.empty([3, 4])
c = a[:, idx[:,:,0], idx[:,:,1]]
for i in range(len(b)):
    b[i] = c[i][i]
    
print(b)
#[[ 0.  1.  2.  3.]
#[40. 41. 42. 43.]
#[80. 81. 82. 83.]]

Answer 3

b = np.diagonal(a[:, idx[..., 0], idx[..., 1]]).T
print(b)

# [[ 0  1  2  3]
#  [40 41 42 43]
#  [80 81 82 83]]

Comments

[hirohi_71]

ありがとうございます