Векторизация многократного вызова массива с различными индексами

Question

У меня есть 2D-массив numpy и список списков индексов, для которых я хочу вычислить сумму соответствующих 1D-векторов из массива numpy. Это можно легко сделать с помощью цикла for или через понимание списка, но мне интересно, можно ли его векторизовать. С подобным кодом я получаю примерно 40-кратное ускорение от векторизации.

Вот пример кода:

import numpy as np 
indices = [[1,2],[1,3],[2,0,3],[1]] 
array_2d = np.array([[0.5, 1.5],[1.5,2.5],[2.5,3.5],[3.5,4.5]]) 
soln = [np.sum(array_2d[x], axis=-1) for x in indices] 

(изменить): Обратите внимание, что индексы не являются (х, у) координаты для array_2d, вместо индексов [0] = [1,2] представляет собой первый и второй векторы (строки) в array_2d. Количество элементов каждого списка в индексах может быть переменным.

Это то, что я хотел бы надеяться, чтобы быть в состоянии сделать:

vectorized_soln = np.sum(array_2d[indices[:]], axis=-1) 

Кто-нибудь знает, есть ли какие-нибудь способы достижения этой цели?

Answer 1

Во-первых, чтобы все, я думаю, что у вас есть опечатка в третьем элементе индексов ...

Самый простой способ сделать это, строит sub_array с двумя массивами индексов:

i = np.array([1,1,2]) 
j = np.array([2,3,?]) 
sub_arr2d = array_2d[i,j] 

и, наконец, вы можете взять сумму sub_arr2d ...