如何索引numpy數組？

Question

一個簡單的例子：我有一個numpy數組（[[5,4,2]，[1,8,8]，[4,5,6]]）和另一個數組（[2,2,1]） 。第二個數組包含第一個數組中每行所關注的索引。我想我的最終數組只是一個矢量數組（[2,8,5]）。我正在尋找一些沿着first_array [second_array]的線路，但它證明難以捉摸。謝謝！

Answer 1

你可以嘗試像以下：

>>> arr 
array([[5, 4, 2], 
     [1, 8, 8], 
     [4, 5, 6]]) 
>>> idx 
array([2, 2, 1]) 
>>> arr[np.arange(arr.shape[0]), idx] 
array([2, 8, 5])