如何高效计算Pandas DataFrame中跨趋势迭代的EMA？

问题描述 #

我们有如下Pandas DataFrame：

import pandas as pd
import numpy as np

ds = { 'trend' : [1,1,1,1,2,2,3,3,3,3,3,3,4,4,4,4,4], 'price' : [23,43,56,21,43,55,54,32,9,12,11,12,23,3,2,1,1]}
df = pd.DataFrame(data=ds)

需要新增字段ema2，要求：

• 遍历每一条记录，结合当前行价格及之前各趋势的最后一个价格，计算周期为2的指数移动平均（EMA）
• EMA计算使用com=1（对应周期2），min_periods=2，且采用默认的adjust=True加权方式

原代码通过循环读取前i+1行CSV、分组取最后价格再计算EWM，虽然结果正确，但百万级数据下速度极慢——时间复杂度为O(n²)，且频繁IO操作严重拖慢效率。

优化思路 #

核心优化点：

1. 彻底取消磁盘IO和重复分组操作，直接在内存中维护各趋势的最新价格
2. 手动实现EMA加权计算，避免每次创建Pandas Series的额外开销
3. 采用高效的行遍历方式，将时间复杂度降至O(n*k)（k为趋势数量，远小于n）

优化代码 #

import pandas as pd
import numpy as np

# 初始化数据
ds = { 'trend' : [1,1,1,1,2,2,3,3,3,3,3,3,4,4,4,4,4], 'price' : [23,43,56,21,43,55,54,32,9,12,11,12,23,3,2,1,1]}
df = pd.DataFrame(data=ds)

# EMA参数配置
com = 1  # 对应周期2，com = time_window - 1
alpha = 1 / (1 + com)
ema_list = []

# 维护当前各趋势的最新价格，以及趋势出现的顺序
current_prices = {}
trend_order = []

# 用itertuples遍历，比iterrows效率高2-3倍
for row in df.itertuples(index=False):
    trend = row.trend
    price = row.price
    
    # 更新趋势价格，记录首次出现的趋势顺序
    if trend not in current_prices:
        trend_order.append(trend)
        current_prices[trend] = price
    else:
        current_prices[trend] = price
    
    # 构建当前所有趋势的最新价格序列
    prices = [current_prices[t] for t in trend_order]
    n = len(prices)
    
    if n < 2:
        ema_val = np.nan
    else:
        # 手动计算adjust=True的EWM加权均值
        weight_sum = 0.0
        weighted_sum = 0.0
        # 从最新价格开始计算权重（权重随时间指数衰减）
        for idx, p in enumerate(reversed(prices)):
            weight = (1 - alpha) ** idx
            weight_sum += weight
            weighted_sum += p * weight
        ema_val = weighted_sum / weight_sum
    
    ema_list.append(ema_val)

# 将结果赋值给DataFrame
df['ema2'] = ema_list
print(df)

结果验证 #

运行上述代码后，输出结果与原代码完全一致：

trend  price       ema2
0       1     23        NaN
1       1     43        NaN
2       1     56        NaN
3       1     21        NaN
4       2     43  35.666667
5       2     55  43.666667
6       3     54  49.571429
7       3     32  37.000000
8       3      9  23.857143
9       3     12  25.571429
10      3     11  25.000000
11      3     12  25.571429
12      4     23  24.200000
13      4      3  13.533333
14      4      2  13.000000
15      4      1  12.466667
16      4      1  12.466667

性能提升说明 #

• 完全取消磁盘IO操作，所有数据在内存中处理，避免了文件读写的巨大开销
• 采用itertuples遍历行数据，比原生iterrows效率提升2-3倍
• 手动实现EMA计算，避免了Pandas Series对象的频繁创建与销毁，计算效率显著提升
• 时间复杂度从原代码的O(n²)降至O(n*k)（k为趋势数量），百万级数据下运行时间可从小时级缩短至秒级

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