如何在Pandas中计算“距离最近高点的K线数”并随新高重置

问题描述 #

我正在尝试在Pandas中计算一个滚动的“距离最近高点的K线数（bars since high）”，该数值会在出现新高时重置。目前我能够计算滚动高点，但无法统计距离该高点的行数。

示例代码：

import pandas as pd

df = pd.DataFrame([0,1,2,3,10,3,4,5,25],columns=['price'])
df['high'] = df['price'].rolling(window=100000,min_periods=1).max()

期望输出：

df['barssincehigh'] = [0,0,0,0,0,1,2,3,0]

解决方案 #

要实现需求，核心是先识别新高点位置，再对每个新高点后的行连续计数，直到下一个新高点出现时重置。具体步骤如下：

• 标记新高点：判断当前价格是否等于滚动高点，生成布尔列标记新高点

  df['is_new_high'] = df['price'] == df['high']
  

• 创建分组标签：对布尔列累计求和，每次新高点出现时分组标签递增，将每个新高点到下一个新高点的行归为同一组

  df['group'] = df['is_new_high'].cumsum()
  

• 组内计数：在每个分组内从0开始累计计数，得到距离最近高点的K线数

  df['barssincehigh'] = df.groupby('group').cumcount()
  

整合后的完整代码：

import pandas as pd

df = pd.DataFrame([0,1,2,3,10,3,4,5,25],columns=['price'])
df['high'] = df['price'].rolling(window=100000,min_periods=1).max()
df['is_new_high'] = df['price'] == df['high']
df['group'] = df['is_new_high'].cumsum()
df['barssincehigh'] = df.groupby('group').cumcount()

# 查看结果
print(df[['price', 'high', 'barssincehigh']])

运行后输出的barssincehigh列与期望一致：

price  high  barssincehigh
0      0     0              0
1      1     1              0
2      2     2              0
3      3     3              0
4     10    10              0
5      3    10              1
6      4    10              2
7      5    10              3
8     25    25              0

如果不需要中间列，可简化为链式调用或删除中间列：

# 链式调用写法
df['barssincehigh'] = df.groupby((df['price'] == df['high']).cumsum()).cumcount()

# 删除中间列
df.drop(['is_new_high', 'group'], axis=1, inplace=True)

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