Vectorize Loop Using Numpy

Question

У меня есть два массива 1D Numpy index_data и charge, которые имеют одинаковую длину и содержат int и floats соответственно. Я пытаюсь сделать total_charge массив следующим образом:

total_charge = [] 
for i in range(len(index_data)): 
    if index_data[i] == 0: 
     total_charge.append(charge[i]) 
    else: 
     total_charge[-1] += charge[i] 
total_charge = np.array(total_charge) 

Как бы я векторизации это? Помоги мне, колдуны, твоя единственная надежда.

Answer 1

Конкретный пример index_data и charge поможет уточнить ожидаемый ввод. Однако, прочитав то, что у вас есть, я думаю, что можно использовать np.add.reduceat.

По-видимому, каждая новая группа сборов обозначается 0 в index_data; в противном случае значения в index_data не используются. total_charge - сумма сборов в каждой группе.

Вот пример того, как можно использовать reduceat.

Во-первых, некоторые демо-данные:

In [105]: index_data 
Out[105]: array([0, 1, 1, 0, 1, 1, 1, 0, 1]) 

In [106]: charge 
Out[106]: array([ 1.5, 2. , 3. , 2.5, 1.5, 1. , 1. , 2. , 1. ]) 

zerolocs дает индексы, где 0 происходит в index_data:

In [107]: zerolocs = where(index_data==0)[0] 

In [108]: zerolocs 
Out[108]: array([0, 3, 7]) 

Использование np.add.reduceat просуммировать группы зарядов.

In [109]: total_charge = np.add.reduceat(charge, zerolocs) 

In [110]: total_charge 
Out[110]: array([ 6.5, 6. , 3. ])