Pandas向量化操作：实现阈值达5时重置的计数功能

可以通过Pandas的分组与累计计数实现向量化操作，无需循环，步骤如下：

1. 标记信号变化的位置：当当前行sign与上一行不同时，标记为需要重置计数的节点。
2. 计算无阈值限制的连续计数，标记达到阈值5的行（其下一行需重置计数）。
3. 合并两种重置条件生成分组键，在每个分组内重新累计计数。

完整代码示例：

import pandas as pd

# 示例数据
data = {
    'price': [13,12,11,12,13,14,14,14,14,14,14],
    'sign': [1,1,-1,-1,1,1,1,1,1,1,1]
}
df = pd.DataFrame(data)

# 标记sign变化的行
df['sign_change'] = df['sign'] != df['sign'].shift()

# 计算初始连续计数（不考虑阈值），标记前一行计数到5的位置
df['temp_count'] = df.groupby(df['sign_change'].cumsum()).cumcount() + 1
df['need_reset'] = df['temp_count'].shift().eq(5)

# 合并重置条件生成分组键
df['group_key'] = (df['sign_change'] | df['need_reset']).cumsum()

# 在每个分组内计算最终count
df['count'] = df.groupby('group_key').cumcount() + 1

# 清理中间辅助列（可选）
df = df.drop(['sign_change', 'temp_count', 'need_reset', 'group_key'], axis=1)

print(df)

运行后输出结果与期望一致：

price  sign  count
0      13     1      1
1      12     1      2
2      11    -1      1
3      12    -1      2
4      13     1      1
5      14     1      2
6      14     1      3
7      14     1      4
8      14     1      5
9      14     1      1
10     14     1      2

原理说明 #

• sign_change捕捉信号切换的节点，作为第一个重置条件。
• temp_count先计算常规连续信号计数，need_reset标记前一行计数达5的位置，作为第二个重置条件。
• group_key将每一段需要连续计数的区间划分为独立分组，每个分组内从1开始累计计数，完全匹配需求。

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