Индексирование массива 2D numpy массивом индексов столбцов

Question

Скажем, у меня есть двумерный массив numpy A формы (N, D) и одномерный массив b формы (N,). Я хочу проиндексировать первый массив, используя второй, чтобы получить следующий массив:

arr[0] = A[0][b[0]]
arr[1] = A[1][b[1]]
arr[2] = A[2][b[2]]
...
arr[N-1] = A[N-1][b[N-1]]

Это легко сделать с помощью понимания списка, но я надеюсь на эффективное векторизованное решение с методами numpy или индексации.

Answer 1

Попробуйте это:

for i,x in enumerate(a): #loop iterate from a[0] to a[N-1]
  arr[i] = x[b[i]] #b[i] is from b[0] to b[N-1]

Answer 2

Вы можете использовать numpy индексирование целых массивов

arr = A[range(N), b]