numpy配列を効果的にスライスする方法

Question

forループではなく、行列演算で次のコードを実装します。

基本的には、第2次元で最大値のインデックスを取得したいと考えています。そして、私はこれらの指数に従ってbの対応する3次元値を抽出したいと思います。

たとえば：

その後、 np.argmax(a, axis=1)

a: 
array([[5, 4, 1], 
     [3, 1, 3], 
     [4, 1, 2], 
     [0, 0, 5]]) 
b: 
array([[[1, 3], [1, 4], [5, 0]], 
     [[2, 4], [2, 2], [1, 2]], 
     [[2, 1], [1, 2], [4, 5]], 
     [[4, 0], [5, 5], [0, 2]]]) 

はarray([0, 0, 0, 2]) だから

c[0] = b[0][0], c[1]=b[1]b[0], c[2]=b[2][0], c[3]=b[3][2]は、私は、これは走行速度を遅くするループのために、より高速でこれを達成するために任意のよりエレガントな方法があると思い与えます方法？

Answer 1

あなたはそうのような派手なインデックスを使用することができます。

b[np.arange(4), np.argmax(a, axis=-1)] 
# array([[1, 3], 
#  [2, 4], 
#  [2, 1], 
#  [0, 2]])