如何用Pandas实现基于indicator列的条件累加并生成period列

原始数据构造 #

import numpy as np
import pandas as pd

indicator = ["buy"] + ["hold"]*3 + ["sell"] + ["hold"]*4 + ["buy"] + ["hold"] * 2
values = np.random.randn(len(indicator)) / 100
df = pd.DataFrame({"indicator": indicator, "values": values})

原始输出：

indicator   values
0   buy      0.001810
1   hold     0.011779
2   hold    -0.003350
3   hold    0.010311
4   sell    -0.010846
5   hold    -0.013635
6   hold    0.003794
7   hold    -0.003792
8   hold    0.006421
9   buy    -0.019779
10  hold    0.007123
11  hold    0.025983

需求说明 #

• 当出现buy时，累加后续所有行直到遇到sell
• 当出现sell时，累加后续所有行直到遇到buy
• 标记每个阶段的周期编号（period列），最终输出包含period和累计值Cumsum的DataFrame

解决方案 #

1. 生成period列 #

通过识别buy/sell作为周期起始点，对起始点的标记做累加得到周期编号：

# 标记所有周期起始行（buy或sell）
start_points = df['indicator'].isin(['buy', 'sell'])
# 生成period列
df['period'] = start_points.cumsum()

2. 计算分段累加值 #

用groupby按周期分组后，对values做累计求和：

df['Cumsum'] = df.groupby('period')['values'].cumsum()

完整代码与输出 #

import numpy as np
import pandas as pd

# 构造原始DataFrame
indicator = ["buy"] + ["hold"]*3 + ["sell"] + ["hold"]*4 + ["buy"] + ["hold"] * 2
values = np.random.randn(len(indicator)) / 100
df = pd.DataFrame({"indicator": indicator, "values": values})

# 生成period列
start_points = df['indicator'].isin(['buy', 'sell'])
df['period'] = start_points.cumsum()

# 计算分段累计值
df['Cumsum'] = df.groupby('period')['values'].cumsum()

print(df)

预期输出：

indicator   values  period  Cumsum
0   buy     0.004730    1   0.004730
1   hold    -0.006814   1   -0.002084
2   hold    0.002424    1   0.000340
3   hold    -0.017007   1   -0.016667
4   sell    0.007531    2   0.007531
5   hold    -0.015347   2   -0.007816
6   hold    0.000051    2   -0.007765
7   hold    -0.001202   2   -0.008967
8   hold    -0.008070   2   -0.017037
9   buy     0.028718    3   0.028718
10  hold    -0.005978   3   0.022740
11  hold    0.004725    3   0.027465

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