基于Pandas实现Yes到下一个Close间Returns的求和计算

思路 #

核心是通过分组标记+向前填充识别每个Yes到对应下一个Close的区间，对区间内的Return求和后，仅在每个区间的第一个Yes行写入总和，避免逐行遍历的繁琐操作。

完整代码 #

import pandas as pd

# 创建示例数据集
df = pd.DataFrame(
    {'Strategy':[0,'Close','Yes',0,'Close','Close',0,'Close','Yes','Yes','0','Close'], 
     'Return':[200,200,400,400,400,500,500,500,300,300,500,300]}
)

# 统一Strategy列类型，将数值0转为字符串'0'，避免类型判断冲突
df['Strategy'] = df['Strategy'].astype(str)

# 1. 创建分组分界点：将Close的标记向下偏移一行，确保Yes到Close的所有行在同一分组
df['split'] = df['Strategy'].eq('Close').shift(fill_value=False).cumsum()

# 2. 标记区间起始Yes的索引：仅Yes行保留自身索引，其余为NA，再在分组内向前填充
df['yes_group'] = df.index.where(df['Strategy'].eq('Yes')).groupby(df['split']).ffill()

# 3. 计算每个区间的Return总和
total_returns_map = df.groupby('yes_group')['Return'].sum()

# 4. 创建Total Returns列，仅在每个区间的第一个Yes行写入总和，其余为NA
df['Total Returns'] = pd.NA
first_yes_indices = df[df['Strategy'].eq('Yes')].groupby('yes_group').head(1).index
df.loc[first_yes_indices, 'Total Returns'] = df.loc[first_yes_indices, 'yes_group'].map(total_returns_map)

# 清理临时列（可选）
df = df.drop(['split', 'yes_group'], axis=1)

print(df)

代码解释 #

1. 统一列类型：将Strategy列的数值0转为字符串'0'，避免后续字符串匹配时的类型不一致问题。
2. 分组分界点：通过shift把Close的标记下移一行，确保Yes到对应Close的所有行（包括Close本身）被划分到同一分组。
3. 区间标记：用index.where仅在Yes行保留自身索引，其余行设为NA；再通过groupby+ffill将该索引填充到整个区间的所有行，实现对完整区间的标记。
4. 区间求和：按区间标记分组计算Return总和，得到每个Yes对应的区间总收益。
5. 写入结果：仅在每个区间的第一个Yes行写入总和，其余行保持NA，满足“以第一个Yes为准”的要求。

输出结果 #

示例数据集执行后，Total Returns列会在索引2（第一个Yes）显示1200，索引8（最后一组第一个Yes）显示1100，其余行为NA，完全匹配预期结果。

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