如何向量化处理计算上次条件满足后K线数的for循环？

需求说明 #

给定包含NaN、0、1的序列，其中0为计数锚点（代表条件满足），需要生成如下结果：

• 0位置保持为0
• 后续位置依次递增计数（1、2、3...），直到遇到下一个0重置为0
• 原始NaN位置保持NaN不变

示例：

• 初始序列：NaN, NaN, NaN, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 0
• 期望输出：NaN, NaN, NaN, 0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3, 4, 5, 0

原循环代码的问题 #

你当前的for循环存在以下问题：

• 语法错误：if cond_count['Bars Since'].iloc[i] == 'NaN' 缺少结尾冒号
• 变量j未初始化，while循环无法正常执行
• 逐行修改iloc的方式性能极低，数据量大时会严重拖慢速度

向量化解决方案 #

直接用pandas的分组累加实现，完全替代循环：

import pandas as pd

# 构造示例序列（如果是DataFrame列，替换为df['你的列名']即可）
s = pd.Series([pd.NA, pd.NA, pd.NA, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 0])

# 生成分组标记：NaN块单独分组，每个0到下一个0的区域为一组
groups = s.isna().cumsum() + (s == 0).cumsum()

# 组内累计计数，0位置刚好是组内第一个元素（计数0）
result = s.groupby(groups).cumcount()

# 还原原始NaN位置的缺失值
result[s.isna()] = pd.NA

# 查看结果
print(result.tolist())
# 输出：[<NA>, <NA>, <NA>, 0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3, 4, 5, 0]

逻辑解释 #

1. 分组标记：s.isna().cumsum() 给连续的NaN块分配唯一编号；(s == 0).cumsum() 每次遇到0就增加分组编号，确保每个0到下一个0的区间属于独立分组。
2. 组内计数：groupby(groups).cumcount() 对每个组内的元素从0开始计数，刚好匹配需求：0位置是组内第一个元素（计数0），后续元素依次递增。
3. 还原缺失值：最后将原始序列中的NaN位置对应结果设为NaN，保持数据一致性。

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