如何向量化优化Pandas带条件逐行循环 解决大规模数据迭代过慢问题

原代码耗时久的核心原因：逐行遍历8万行员工数据时，每一行都对排班表做一次全表扫描过滤，时间复杂度为O(员工行数 * 排班表行数)，数据量越大耗时呈指数级上升。

核心优化逻辑 #

完全抛弃逐行遍历，改用「先聚合排班表 + 批量左连接」的向量化方案，所有计算都是Pandas底层C实现，8万行规模下耗时会从数小时降到秒级，且业务逻辑与原代码完全一致。

分步实现说明 #

1. 构造联合匹配键 #

给两个表生成Building + ID的联合匹配字段，后续匹配直接用这个键，避免每行都判断两个条件：

employee_df['Building_ID'] = employee_df['Building'].astype(str) + '_' + employee_df['ID'].astype(str)
schedule['Building_ID'] = schedule['Building'].astype(str) + '_' + schedule['ID'].astype(str)

2. 一次性聚合排班表 #

按联合键分组，一次计算完所有需要的统计值，不用每行员工数据都重复计算：

schedule_agg = schedule.groupby('Building_ID').agg(
    SumDuration=('Dur', 'sum'),
    MinDuration=('Dur', 'min'),
    MaxDuration=('Dur', 'max'),
    Sequence=('WorkType', lambda x: ''.join(x.tolist()))  # 直接拼接同组所有WorkType
).reset_index()

3. 批量左连接匹配 #

左连接的逻辑天然匹配原代码「无匹配条目就保留员工原始数据、填充NA」的规则，不用写额外的判断分支：

merged = employee_df.merge(schedule_agg, on='Building_ID', how='left')

4. 向量化计算剩余字段 #

批量处理时长转换、百分比计算，完全替代原代码的逐行时间格式解析：

# 批量把HH:MM:SS格式的Duration转成秒数
merged['Duration_seconds'] = pd.to_timedelta(merged['Duration']).dt.total_seconds()
# 批量计算%Time
merged['%Time'] = 1.0 / merged['Duration_seconds']
# Gap列原逻辑固定填NA，直接生成列
merged['Gap'] = 'NA'

5. 整理输出结构 #

按原输出列顺序整理，空值自动填充为NA，完全匹配原输出格式：

output_columns = ['Name', 'Date', 'Building', 'ID', 'Years', 'Type', 'Start Time', 'Stop Time', 'Duration',
                  'SumDuration', '%Time', 'Gap', 'Sequence', 'MinDuration', 'MaxDuration', 'Status']
output_df = merged[[
    'Name', 'Date', 'Building', 'ID', 'Years', 'EmployeeType',
    'Start Time', 'Stop Time', 'Duration', 'SumDuration', '%Time',
    'Gap', 'Sequence', 'MinDuration', 'MaxDuration', 'Status'
]].rename(columns={'EmployeeType': 'Type'}).fillna('NA')

优化后的完整可运行代码 #

import pandas as pd

pd.set_option('display.expand_frame_repr', False)
pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)


def create_output(employee_file, schedule_file, output_file):
    # 读取时可指定dtypes减少内存占用，进一步提速
    employee_df = pd.read_csv(employee_file)
    schedule = pd.read_csv(schedule_file)

    # 构造联合匹配键
    employee_df['Building_ID'] = employee_df['Building'].astype(str) + '_' + employee_df['ID'].astype(str)
    schedule['Building_ID'] = schedule['Building'].astype(str) + '_' + schedule['ID'].astype(str)

    # 排班表一次性聚合
    schedule_agg = schedule.groupby('Building_ID').agg(
        SumDuration=('Dur', 'sum'),
        MinDuration=('Dur', 'min'),
        MaxDuration=('Dur', 'max'),
        Sequence=('WorkType', lambda x: ''.join(x.tolist()))
    ).reset_index()

    # 左连接匹配
    merged = employee_df.merge(schedule_agg, on='Building_ID', how='left')

    # 向量化计算字段
    merged['Duration_seconds'] = pd.to_timedelta(merged['Duration']).dt.total_seconds()
    merged['%Time'] = 1.0 / merged['Duration_seconds']
    merged['Gap'] = 'NA'

    # 整理输出列
    output_columns = ['Name', 'Date', 'Building', 'ID', 'Years', 'Type', 'Start Time', 'Stop Time', 'Duration',
                      'SumDuration', '%Time', 'Gap', 'Sequence', 'MinDuration', 'MaxDuration', 'Status']
    output_df = merged[[
        'Name', 'Date', 'Building', 'ID', 'Years', 'EmployeeType',
        'Start Time', 'Stop Time', 'Duration', 'SumDuration', '%Time',
        'Gap', 'Sequence', 'MinDuration', 'MaxDuration', 'Status'
    ]].rename(columns={'EmployeeType': 'Type'}).fillna('NA')

    # 后续业务计算逻辑不变
    output_df.to_csv(output_file, index=False)


def main():
    create_output('employees.csv', 'schedule.csv', 'out.csv')


if __name__ == '__main__':
    main()

额外提速技巧 #

• 读取CSV时通过dtypes参数指定数值列的类型（比如ID、Building设为int32，Dur设为float32），减少内存占用30%以上，读取速度同步提升
• 去掉原代码循环里的print语句，大量IO操作会严重拖慢执行速度
• 若排班表的联合键无重复值，可把聚合后的结果转成字典做映射，速度比merge更快

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