索引一個索引數組的多維numpy數組

Question

我很難索引一個索引數組的二維數組。

假設我有：

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

# Array to be indexed 
my_array = np.array([[0,2,1], [0,5,1], [1,2,1], [5,1,3], [2,6,2]]) 

和預期輸出應該是0指數my_array第一項的my_array第二項的2指數等，因此：

# Expected output 
expected_array = np.array([0,1,2,3,6]) 

我已經使用for循環實際上做了它：

# Using for-loop 
expected_array = np.array([]) 
for i in range(len(my_array)): 
    expected_array = np.append(idx, my_array[i][idx[i]]) 

但我想知道是否有辦法可能做陣列索引不使用for-loops？假設len(idx)==len(my_array)一直在。

Answer 1

您可以使用numpy.arange來創建一個數組來索引行，這些行一次只是每行。然後使用idx數組來索引每行。

>>> my_array[np.arange(my_array.shape[0]), idx] 
array([0, 1, 2, 3, 6])