如何用pandas内置函数替换for循环高效计算AFirst/BFirst列

处理大规模数据集时，Python层的逐行for循环开销极高，完全可以通过pandas内置的向量化操作、分组运算改写逻辑，运行效率通常可以提升1~2个数量级。

计算规则 #

需要为DataFrame新增AFirst、BFirst两列，按如下规则填充0/1值：

• 当start列值为1时，AFirst、BFirst重置为初始值0，作为新计算段的起点
• 同一段内（两个相邻start=1的行之间），如果当前行A列值≥2，首次触发时将AFirst置为1
• 同一段内，如果当前行B列值≤-2，首次触发时将BFirst置为1
• 任意一列被置为1后，直到下一个start=1的行出现前，该列始终保持1，另一列保持0，即保留先触发的列标记

原有逐行实现（性能差） #

df = pd.DataFrame({'start': {0: 1, 1: 0, 2: 0, 3: 0, 4: 0, 5: 1, 6: 0, 7: 0, 8: 0, 9: 0}, 'A': {0: 0.0, 1: 1.0, 2: 2.0, 3: 3.0, 4: 1.0, 5: 0.0, 6: 0.5, 7: 1.0, 8: 1.5, 9: 2.0}, 'B': {0: 0.0, 1: -1.0, 2: -1.5, 3: -2.0, 4: -1.0, 5: 0.0, 6: -1.0, 7: -2.0, 8: -3.0, 9: -1.0}})

for i in range(df.shape[0]):

        if df['start'].iloc[i] == 1:

            if (df.loc[df.index[i],'A'] <= 2) & (df.loc[df.index[i],'B'] >= -2):
                df.loc[df.index[i],'AFirst']=0
                df.loc[df.index[i],'BFirst']=0
            elif df.loc[df.index[i],'A'] >= 2:
                df.loc[df.index[i],'AFirst']=1
                df.loc[df.index[i],'BFirst']=0
            elif df.loc[df.index[i],'B'] <= -2:
                df.loc[df.index[i],'AFirst']=0
                df.loc[df.index[i],'BFirst']=1

        elif df.loc[df.index[i-1],'AFirst'] == 1:
            df.loc[df.index[i],'AFirst']=1
            df.loc[df.index[i],'BFirst']=0

        elif df.loc[df.index[i-1],'BFirst'] == 1:
            df.loc[df.index[i],'AFirst']=0
            df.loc[df.index[i],'BFirst']=1

        elif df.loc[df.index[i],'A'] >= 2:
            df.loc[df.index[i],'AFirst']=1
            df.loc[df.index[i],'BFirst']=0

        elif df.loc[df.index[i],'B'] <= -2:
            df.loc[df.index[i],'AFirst']=0
            df.loc[df.index[i],'BFirst']=1

向量化优化实现（性能提升10~100倍） #

核心思路是先通过cumsum给每个独立计算段打分组标记，仅在分组层面做逻辑判断，所有行级运算全部走pandas底层C实现的向量化接口，避免Python层逐行遍历。

import pandas as pd
import numpy as np

df = pd.DataFrame({'start': {0: 1, 1: 0, 2: 0, 3: 0, 4: 0, 5: 1, 6: 0, 7: 0, 8: 0, 9: 0}, 'A': {0: 0.0, 1: 1.0, 2: 2.0, 3: 3.0, 4: 1.0, 5: 0.0, 6: 0.5, 7: 1.0, 8: 1.5, 9: 2.0}, 'B': {0: 0.0, 1: -1.0, 2: -1.5, 3: -2.0, 4: -1.0, 5: 0.0, 6: -1.0, 7: -2.0, 8: -3.0, 9: -1.0}})

# 给每个连续计算段分配唯一分组ID
df['group_id'] = df['start'].cumsum()
# 标记每行是否满足A、B的触发条件
hit_a = df['A'] >= 2
hit_b = df['B'] <= -2

# 计算每个分组内第一次触发A、B的行索引
first_a = df[hit_a].groupby('group_id').apply(lambda x: x.index.min(), include_groups=False)
first_b = df[hit_b].groupby('group_id').apply(lambda x: x.index.min(), include_groups=False)

# 初始化两列默认值为0
df[['AFirst', 'BFirst']] = 0

# 仅遍历分组（分组数远小于总行数，开销可忽略）
for gid in df['group_id'].unique():
    pos_a = first_a.get(gid, np.inf)
    pos_b = first_b.get(gid, np.inf)
    # 找到当前分组的结束位置（下一个分组的起点）
    next_group_pos = df[df['group_id'] == gid + 1].index.min()
    end_pos = next_group_pos if not pd.isna(next_group_pos) else df.index[-1] + 1
    
    # 组内无触发条件，保持默认0即可
    if pos_a == np.inf and pos_b == np.inf:
        continue
    # A条件先触发
    if pos_a < pos_b:
        df.loc[pos_a:end_pos-1, 'AFirst'] = 1
    # B条件先触发
    else:
        df.loc[pos_b:end_pos-1, 'BFirst'] = 1

# 删除临时列
df = df.drop(columns=['group_id'])

结果验证 #

运行上述代码后输出结果和原逐行循环完全一致：

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