Получение значений из массива numpy с использованием массива индексов ⇐ Python

[Anonymous]

Я пытаюсь найти простую команду numpy для получения значений трехмерной матрицы в соответствии с индексным массивом для второй координаты. Это можно сделать следующим образом:

Код: Выделить всё

import numpy as np
entry = np.array([[1, 2],[3,4],[5, 6]])
entries = np.stack([entry, entry, entry, entry])
indices = np.array([2, 1, 0, 1])

r = np.array([entries[i, indices[i], :] for i in range(len(indices))])

Этот код дает правильный результат, но использует понимание списка вместо одной команды.
Есть ли numpy-операция, которая делает то же самое? Я ожидал, что np.take() будет работать, но это не так, поскольку он создает массив формы (4, 4, 2) вместо (4, 2)

Подробнее здесь: https://stackoverflow.com/questions/792 ... of-indices