如何高效对多时段DataFrame列表按类型计算向量求和与乘积

问题背景 #

拥有对应0到N不同时段的DataFrame列表，每个DataFrame包含多种类型，需要针对每种类型高效计算指定的求和（∑）与乘积（∏）运算，公式如下：

示例数据集 #

为简化演示，列表中每个DataFrame设为相同值，计算逻辑不变：

import pandas as pd

d = {'type': ['a', 'b', 'c'], 'x': [1, 2, 3], 'y':[3,4,5]}
df = pd.DataFrame(data=d)

l = []
window = 20
[l.append(df.copy(deep=True)) for i in range(window)]

低效实现方式 #

以下实现方式在大数据集/大窗口下效率极低，核心问题在于嵌套循环和多次切片过滤：

import pandas as pd
import math

d = {'type': ['a', 'b', 'c'], 'x': [1, 2, 3], 'y':[3, 4, 5]}
df = pd.DataFrame(data=d)

l = []
window = 20
for i in range(window):
    df['window'] = i
    l.append(df.copy(deep=True))
df = pd.concat(l)

sums = {}
types = df['type'].unique()
for s in types:
    sums[s] = 0
    tdf = df[df['type'] == s]
    for i in range(window):
        sums[s] += tdf[tdf['window'] == i]['x'].values[0] * math.prod(2 - tdf[tdf['window'] == i+1]['x'].values[0] for i in range(0, window-1)) * tdf[tdf['window'] == i]['y'].values[0]

高效实现方案 #

利用Pandas的向量化运算和分组累积操作，完全避免嵌套循环，大幅提升计算效率：

核心思路 #

1. 合并所有时段的DataFrame并添加时段标识，确保数据按type和window有序排列
2. 按type分组后，对每个组计算(2 - x)的序列
3. 从后往前计算累积乘积，得到每个时段i对应的∏_{k=i+1}^{N-1}(2-x_k)值（最后一个时段的乘积为1，因为无后续元素）
4. 计算每个时段的x_i * y_i * 累积乘积，再求和得到该类型的最终结果

代码实现 #

import pandas as pd

d = {'type': ['a', 'b', 'c'], 'x': [1, 2, 3], 'y':[3,4,5]}
df = pd.DataFrame(data=d)

window = 20
# 合并所有时段数据并添加window列
full_df = pd.concat([df.assign(window=i) for i in range(window)], ignore_index=True)
# 确保按type和window排序，保证时段顺序正确
full_df = full_df.sort_values(['type', 'window']).reset_index(drop=True)

def calculate_type_result(group):
    # 计算(2 - x)的序列
    term = 2 - group['x']
    # 从后往前计算累积乘积，最后一个元素设为1（空乘积的结果）
    cum_prod = term.iloc[::-1].cumprod().iloc[::-1]
    # 处理最后一个时段：后面没有元素，乘积为1
    cum_prod = cum_prod.shift(-1).fillna(1)
    # 计算每个时段的贡献并求和
    return (group['x'] * group['y'] * cum_prod).sum()

# 按type分组计算，得到最终结果
result = full_df.groupby('type').apply(calculate_type_result).to_dict()
print(result)

效率说明 #

• 所有操作均为Pandas内置的向量化运算，避免了Python层面的循环，计算速度比原实现提升几个数量级
• 分组操作仅执行一次，无需反复过滤数据
• 累积乘积操作利用Pandas的cumprod方法，底层为C语言实现，性能远超math.prod加循环的组合

内容的提问来源于stack exchange，提问作者pyCthon