fancy indexing numpy matrix: один элемент в строке

Question

У меня есть матрица размером 2d, матрица, формы (m, n). Мой фактическое использование регистр имеет т ~ 1e5 и п ~ 100, но ради иметь простой минимальный пример:

matrix = np.arange(5*3).reshape((5, 3)) 

У меня есть индексация массив целых чисел, IDX, по форме (м,), с каждая запись между [0, n). Этот массив указывает, какой столбец следует выбирать из каждой строки матрицы .

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

Итак, я пытаюсь выбрать столбец 2 из ряда 0, столбец 0 из строки 1, столбец 2 из строки 2, столбец 1 из строки 1 и столбца 1 из строки 4. Таким образом, окончательный ответ должен быть:

correct_result = np.array((2, 3, 8, 10, 13)) 

Я попробовал следующее, которое интуитивно, но неправильно:

incorrect_result = matrix[:, idx] 

Что выше синтаксис делает применить IDX как фантазии массива индексации, строка за строкой, в результате чего в другой матрица формы (m, n), что я не хочу.

Каков правильный синтаксис для фантазийного индексирования этого типа?

Answer 1

correct_result = matrix[np.arange(m), idx] 

Передовое выражение индексирования matrix[I, J] дает выход такому, что output[n] == matrix[I[n], J[n]].

Если мы хотим output[n] == matrix[n, idx[n]], то нам нужно I[n] == n и J[n] == idx[n], так что нам нужно I быть np.arange(m) и J быть idx.