基于cumprod与Period条件批量生成Pandas新列的优化问询

我有一个包含两列的Pandas DataFrame df，分别为df["Period"]和df["Returns"]。df["Period"]的值为1、2、3……n且呈递增状态，每个唯一Period对应的行数无规律且各不相同。

需要基于df["Returns"]的.cumprod()方法，按df["Period"] >= 1、df["Period"] >= 2等条件生成n个新列：

• df["M_1"]：df["Period"] >= 1的行中df["Returns"]的cumprod计算结果
• df["M_2"]：df["Period"] >= 2的行中df["Returns"]的cumprod计算结果
• ……

现有实现方案存在两个缺陷：

1. 当唯一Period数量较多时运行速度极慢
2. 无法很好适配Pandas方法链式调用

示例低效代码：

import numpy as np 
import pandas as pd

# 创建示例数据
n = 10
data = {"Period": np.sort(np.random.randint(1,5,n)),
        "Returns": np.random.randn(n)/100, }
df = pd.DataFrame(data)

# 低效实现
periods = set(df["Period"])
for period in periods:
    cumret = (1 + df.query("Period >= @period")["Returns"]).cumprod() - 1
    df[f"M_{period}"] = cumret
df.head()

预期输出：

方法一：向量化计算，避免循环 #

核心思路是先计算全局的累积乘积，再针对每个Period通过除法推导对应区间的累积收益，同时用掩码处理NaN场景，彻底避免循环中重复的筛选和cumprod计算。

import numpy as np
import pandas as pd

# 创建示例数据
n = 10
data = {"Period": np.sort(np.random.randint(1,5,n)),
        "Returns": np.random.randn(n)/100, }
df = pd.DataFrame(data)

# 计算(1+Returns)的全局累积乘积
df['cum_prod'] = (1 + df['Returns']).cumprod()

# 获取排序后的唯一Period列表
periods = sorted(df['Period'].unique())

# 批量生成M列
for p in periods:
    # 定位第一个符合Period >= p条件的行索引
    first_idx = df[df['Period'] >= p].index[0]
    # 计算基准值：若存在前序行则取前序累积乘积，否则取1
    base = df['cum_prod'].iloc[first_idx - 1] if first_idx > 0 else 1
    # 生成M_p列：符合条件的行用全局累积乘积除以基准值再减1，其余行设为NaN
    df[f'M_{p}'] = np.where(df['Period'] >= p, df['cum_prod'] / base - 1, np.nan)

# 删除临时辅助列
df = df.drop('cum_prod', axis=1)
df.head()

该方案优势：

• 时间复杂度从O(n*m)降至O(n+m)（n为行数，m为唯一Period数量），数据量越大性能提升越明显
• 仅需一次全局累积乘积计算，避免了循环中重复的筛选和cumprod操作

方法二：适配链式调用的向量化实现 #

如果需要整合到Pandas链式调用流程中，可以将逻辑封装为函数，通过pipe和assign方法实现：

def add_m_columns(df):
    # 计算全局累积乘积
    cum_prod = (1 + df['Returns']).cumprod()
    periods = sorted(df['Period'].unique())
    
    # 构建所有M列的键值对字典
    m_cols = {}
    for p in periods:
        first_idx = df[df['Period'] >= p].index[0]
        base = cum_prod.iloc[first_idx - 1] if first_idx > 0 else 1
        m_cols[f'M_{p}'] = np.where(df['Period'] >= p, cum_prod / base - 1, np.nan)
    
    return df.assign(**m_cols)

# 链式调用示例
df = (pd.DataFrame(data)
      .pipe(add_m_columns))

这个实现完全适配链式调用风格，同时保留了向量化的高效性。

原理说明 #

对于任意Period=p，M_p的本质是从第一个Period>=p的行开始到当前行的(1+Returns)乘积减1。而全局累积乘积cum_prod[i]是从第0行到第i行的乘积，因此从first_idx到第i行的乘积可通过cum_prod[i] / cum_prod[first_idx-1]（若first_idx>0）推导得出，最后减1即可得到目标收益值。

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