Pythonic способ получения значений с помощью матрицы индексов?

Question

У меня есть массив n-n-by-3-x-3 A и другой массив n-n-2-x-2 B, который говорит мне, где взять значение в A вкаждое (h, w) местоположение.Обозначим матрицу результатов как C (ч-х-на-3).

Наивно, я могу сделать

C = np.zeros(A.shape)
for i in range(B.shape[0]):
    for j in range(B.shape[1]):
        C[i, j, :] = A[B[i, j, 0], B[i, j, 1], :]

Но я уверен, что есть более эффективные, болеепитон, и быстрее!

Answer 1

Действительно есть:

>>> import numpy as np
>>> 
>>> h, w = 4, 6
>>> 
>>> A = np.random.randint(0, 100, (h, w, 3))
>>> B = np.random.randint(0, h*w, (h, w, 2)) % [h, w]
>>> 
>>> C = np.zeros(A.shape)
>>> for i in range(B.shape[0]):
...     for j in range(B.shape[1]):
...         C[i, j, :] = A[B[i, j, 0], B[i, j, 1], :]
... 
>>> 
>>> C2 = A[(*np.moveaxis(B, 2, 0),)]
>>> np.all(C == C2)
True