索引一个索引数组的多维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])