Расширенное индексирование Numpy с использованием 2D-массива индексов строк без трансляции вывода

У меня есть array ndarray с ndim 3, а некоторые индексы ndarray idxs с ndim 2, которые определяют индексы для первого измерения array . Первое измерение idxs соответствует второму размеру array , то есть idxs.shape[0] == array.shape[1] .

Я хочу получить результат ndarray с ndim 3 и shape (idxs.shape[1], array.shape[1], array.shape[2]) следующим образом:

 for i0 in range(idxs.shape[1]): for i1 in range(array.shape[1]): result[i0, i1] = array[idxs[i1, i0], i1] 

Как я могу получить это более прямо?

Я думал об использовании расширенной индексации, но я не совсем уверен, как это будет выглядеть.

В Theano, следующие работы:

 dim1 = theano.tensor.arange(array.shape[1]) result = array[idxs[dim1], dim1] 

Ваш цикл for :

 out[i, j] == array[idxs[j, i], j] 

То есть, j- й элемент в idxs дает индекс строки в array для i- го элемента в out . Соответствующий набор индексов столбцов в array представляет собой целые числа последовательностей от 0 до idxs.shape[0] - 1 (в данном случае это то же самое, что и array.shape[1] - 1 , но в общем случае не обязательно ).

Поэтому цикл for можно заменить на операцию индексирования одного массива следующим образом:

 def simplified(array, idxs): return array[idxs.T, np.arange(idxs.shape[0])] 

Мы можем проверить правильность и скорость против функций в ответе @ Дивакара:

 m, n = 500, 400 array = np.random.rand(m, n) idxs = np.random.randint(n, size=(n, m)) print(np.allclose(forloop(array, idxs), simplified(array, idxs))) # True %timeit forloop(array, idxs) # 10 loops, best of 3: 101 ms per loop %timeit broadcasted_indexing(array, idxs) # 100 loops, best of 3: 4.1 ms per loop %timeit simplified(array, idxs) # 1000 loops, best of 3: 1.66 ms per loop 

Создайте 2D-сетку индексов, соответствующих индексированию строк: idxs[i1, i0] и используйте массив N x 1 для индексации столбцов. При индексировании в такой array индексы столбцов будут broadcasted в форму индексов строк. Таким образом, у нас будет broadcasted indexing подход, основанный на broadcasted indexing , например,

 # Get 2D grid of row indices corresponding to two nested loops row_idx = idxs[np.arange(array.shape[1])[:,None],np.arange(idxs.shape[1])] # Use column indices alongwith row_idx to index into array. # The column indices would be broadcasted when put as Nx1 array. result = array[row_idx,np.arange(array.shape[1])[:,None]].T 

Обратите внимание, что, как указано в комментариях @ali_m, np.ix_ также можно использовать для создания row_idx , например:

 row_idx = idxs[np.ix_(np.arange(array.shape[1]),np.arange(idxs.shape[1]))] 

Проверка времени выполнения и проверка результатов

Определения функций:

 def broadcasted_indexing(array,idxs): row_idx = idxs[np.arange(array.shape[1])[:,None],np.arange(idxs.shape[1])] return array[row_idx,np.arange(array.shape[1])[:,None]].T def forloop(array,idxs): result = np.zeros((idxs.shape[1],array.shape[1])) for i0 in range(idxs.shape[1]): for i1 in range(array.shape[1]): result[i0, i1] = array[idxs[i1, i0], i1] return result 

Проверка выполнения и проверка результатов:

 In [149]: # Inputs ...: m = 500 ...: n = 400 ...: array = np.random.rand(m,n) ...: idxs = np.random.randint(0,array.shape[1],(n,m)) ...: In [150]: np.allclose(broadcasted_indexing(array,idxs),forloop(array,idxs)) Out[150]: True In [151]: %timeit forloop(array,idxs) 10 loops, best of 3: 136 ms per loop In [152]: %timeit broadcasted_indexing(array,idxs) 100 loops, best of 3: 5.01 ms per loop