Python: Как получить значения многомерного массива в определенных индексированных позициях?

Question

Если у меня есть многомерный массив numpy, например:

a = [[0 1 2]
     [3 4 5]
     [6 7 8]]

Как я могу получить значения в определенных позициях индекса за один шаг? Например, если бы определить пары индексов, например:

indices = [[0,0], [1,1], [2,2]]

Я бы хотел:

a[indices] = [0, 4, 8]

Обратите внимание, что это работает для одномерных массивов (Python: Как получить значения массива в определенных позициях индекса? ), но я не вижу, как заставить это работать в более чем одном измерении. Я использую Python 3,7.

Answer 1

Как и в одномерном ответе, который вы связали, вы можете сделать это элегантно в 2-х измерениях с помощью numpy:

a = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
rows, columns = zip([0, 0], [1, 1], [2, 2])
print(a[rows, columns])

Вывод print будет :

array([0, 4, 8])

Answer 2

Адаптировано из комментария Прасанны

a = [[0 1 2]
     [3 4 5]
     [6 7 8]]
indices = [[0,0], [1,1], [2,2]]

a[indices] = [a[pair[0]][pair[1]] for pair in indices]

Это работает, используя a для каждого l oop, итерируя по каждой паре в массиве индексов, а затем добавляя значения в каждом данном индексе к вашему финальному результат.

pair - каждая пара в массиве индексов и представляет каждую пару индексов в indices. В вашем примере это будет [0,0] на первой итерации, [1,1] на второй и [2,2] на третьей.