Numpy: 2D доступ к массиву с 2D-массив индексов

Question

У меня есть два массива, один представляет собой матрицу из пар индексов,

a = array([[[0,0],[1,1]],[[2,0],[2,1]]], dtype=int) 

, а другая является матрица данных для доступа на эти индексы

b = array([[1,2,3],[4,5,6],[7,8,9]]) 

и я хочу использовать индексы a для получения записей b. Просто делать:

>>> b[a] 

не работает, так как она дает один ряд б для каждой записи в a, т.е.

array([[[[1,2,3], 
     [1,2,3]], 

     [[4,5,6], 
     [4,5,6]]], 


     [[[7,8,9], 
     [1,2,3]], 

     [[7,8,9], 
     [4,5,6]]]]) 

, когда я хотел бы использовать пару индексов в последних осях a чтобы дать два индекса b:

array([[1,5],[7,8]]) 

есть чистый способ сделать это, или мне нужно, чтобы изменить b и объединить столбцы a в соответствующей форме?

В моей реальной проблеме a имеет около 5 миллионов записей, а b - 100 на 100, я бы хотел избежать циклов.

Answer 1

На самом деле, это работает:

b[a[:, :, 0],a[:, :, 1]] 

дает array([[1, 5], [7, 8]]).

Answer 2

В этом случае, это работает более общее решение

tmp = a.reshape(-1,2) 
b[tmp[:,0], tmp[:,1]] 

Answer 3

А всякий раз, когда вы хотите использовать 2D массив индексов формы (п, т) с произвольной большой размерности т, названный аким, чтобы получить доступ к элементам другого 2D массива формы (п, к), названный B:

# array of index offsets to be added to each row of inds 
offset = np.arange(0, inds.size, inds.shape[1]) 

# numpy.take(B, C) "flattens" arrays B and C and selects elements from B based on indices in C 
Result = np.take(B, offset[:,np.newaxis]+inds) 

Вы можете проверить это с помощью, например:

B = 1/(np.arange(n*m).reshape(n,-1) + 1) 
inds = np.random.randint(0,B.shape[1],(B.shape[0],B.shape[1]))