Pandas按连续蜡烛组计算首开与末收的差值并逐行应用

咱先看看手头的原始DataFrame数据（注：原数据最后两行日期重复，我调整为00:07:00避免歧义，实际数据若无需修改可忽略此调整）：

需求很明确：把每个连续蜡烛组（即color相同、run从1开始连续递增的一组）的首根蜡烛open和末根蜡烛close的差值绝对值算出来，再把这个值赋给组内每一行作为新列。

具体实现步骤 #

咱用Pandas来搞定这个事儿，代码如下：

import pandas as pd

# 构造原始数据
data = {
    'date': ['00:01:00', '00:02:00', '00:03:00', '00:04:00', '00:05:00', '00:06:00', '00:07:00'],
    'open': [100, 102, 104, 106, 105, 101, 102],
    'close': [102, 104, 106, 105, 101, 102, 103],
    'color': ['g', 'g', 'g', 'r', 'r', 'g', 'g'],
    'run': [1, 2, 3, 1, 2, 1, 2]
}
df = pd.DataFrame(data)

# 生成分组键：每次run=1时就是新组的开始
df['group_id'] = (df['run'] == 1).cumsum()

# 定义分组计算函数
def calculate_group_diff(group):
    # 获取组内首行的open和末行的close
    first_open = group['open'].iloc[0]
    last_close = group['close'].iloc[-1]
    # 计算绝对值差值
    diff_abs = abs(last_close - first_open)
    # 给组内所有行添加该差值
    group['diff_abs'] = diff_abs
    return group

# 分组应用函数，最后移除临时分组ID列
result_df = df.groupby('group_id').apply(calculate_group_diff).drop('group_id', axis=1)

print(result_df)

运行后的最终结果 #

输出的结果DataFrame如下：

简单说下逻辑：

• 用(df['run'] == 1).cumsum()生成分组ID，刚好匹配“每次run=1开启新组”的规则；
• 分组后对每组取首尾行的关键值计算差值，再把结果广播到组内所有行；
• 最后删掉临时的分组ID列，得到符合需求的结果。

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