Index 2D массив numpy по 2D массиву индексов без петель

Question

Я ищу векторный способ индексирования numpy.array по numpy.array индексов.

Например:

import numpy as np 

a = np.array([[0,3,4], 
       [5,6,0], 
       [0,1,9]]) 

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

Я хочу построить новый массив, таким образом, что каждая строка (я) в этом массиве является строка (я) из массива a, индексируются строки Inds массива (I). Мой желаемый результат:

array([[ 0., 3.], # a[0][:,[0,1]] 
     [ 6., 0.], # a[1][:,[1,2]] 
     [ 0., 9.]]) # a[2][:,[0,2]] 

я могу добиться этого с помощью цикла:

def loop_way(my_array, my_indices): 
    new_array = np.empty(my_indices.shape) 
    for i in xrange(len(my_indices)): 
     new_array[i, :] = my_array[i][:, my_indices[i]] 
    return new_array 

Но я ищу чистый векторизованный раствор.

Answer 1

При использовании массивов индексов для индексации другого массива форма каждого индексного массива должна соответствовать форме вывода . Вы хотите, чтобы индексы столбцов, чтобы соответствовать inds, и вы хотите, индексы строк, чтобы соответствовать строку вывода, что-то вроде:

array([[0, 0], 
     [1, 1], 
     [2, 2]]) 

Вы можете просто использовать один столбец выше, из-за вещания, так что вы можно использовать np.arange(3)[:,None] является вертикальным arange, потому что None вставляет новую ось:

>>> np.arange(3)[:, None] 
array([[0], 
     [1], 
     [2]]) 

Наконец, вместе:

>>> a[np.arange(3)[:,None], inds] 
array([[0, 3], # a[0,[0,1]] 
     [6, 0], # a[1,[1,2]] 
     [0, 9]]) # a[2,[0,2]] 

Answer 2

Это возможно, хотя и несколько неочевидных, чтобы сделать это следующим образом:

>>> a[np.arange(a.shape[0])[:, None], inds] 
array([[0, 3], 
     [6, 0], 
     [0, 9]]) 

Индекс np.arange(a.shape[0]) просто индексируют строки, к которым массив индексов столбцов inds применяется. Добавление [:, None] изменяет форму этого массива таким образом, что его форма равна (a.shape[0], 1), то есть каждый индекс строки находится в отдельной строке двумерного массива шириной 1 столбец.

Основной принцип заключается в том, что количество измерений в массивах индексов должно совпадать, и их формы также должны делать это. См. Документацию для np.ix_, чтобы понять это.