python - 使用索引数组对多维 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 第一个条目的索引，2 my_array 第二个条目的索引依此类推，因此:

# 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 循环的情况下进行数组索引？假设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])