Векторизация выборочной кумулятивной суммы

Question

У меня есть pandas Series, где каждый элемент представляет собой список с индексами:

series_example = pd.Series([[1, 3, 2], [1, 2]])

Кроме того, у меня есть массив со значениями, связанными с каждым индексом:

arr_example = np.array([3., 0.5, 0.25, 0.1])

Я хочу создать новый Series с кумулятивными суммами элементов массива, заданными индексами в строке ввода Series. В этом примере вывод Series будет иметь следующее содержимое:

0    [0.5, 0.6, 0.85]
1    [0.5, 0.75]
dtype: object

Невекторный способ сделать это будет следующим:

def non_vector_transform(series, array):
    series_output = pd.Series(np.zeros(len(series_example)), dtype = object)
    for i in range(len(series)):
        element_list = series[i]
        series_output[i] = []   
        acum = 0
        for element in element_list:
            acum += array[element]
            series_output[i].append(acum)
    return series_output

Я бы хотел сделать это векторизованным способом. Любой волшебник векторизации, чтобы помочь мне здесь?

Answer 1

Используйте Series.apply и np.cumsum:

import numpy as np
import pandas as pd
series_example = pd.Series([[1, 3, 2], [1, 2]])
arr_example = np.array([3., 0.5, 0.25, 0.1])
result = series_example.apply(lambda x: np.cumsum(arr_example[x]))
print(result)

Или, если вы предпочитаете цикл for:

import numpy as np
import pandas as pd
series_example = pd.Series([[1, 3, 2], [1, 2]])
arr_example = np.array([3., 0.5, 0.25, 0.1])
# Copy only if you do not want to overwrite the original series
result = series_example.copy()
for i, x in result.iteritems():
    result[i] = np.cumsum(arr_example[x])
print(result)

Вывод:

0    [0.5, 0.6, 0.85]
1         [0.5, 0.75]
dtype: object