百万行Pandas DataFrame循环提速及代码改写求助

原代码逐行遍历1000万行数据，每次都要切片、过滤分组，时间复杂度是O(n²)，完全没法高效运行。下面用Pandas的向量化分组滚动操作替代循环，速度能提升几个数量级：

import pandas as pd
import numpy as np

# 先把average<=0的值替换为NaN，计算均值时会自动忽略这些数据
data['avg_filtered'] = np.where(data['average'] > 0, data['average'], np.nan)

# 按ACOUNTNO分组，计算当前行之前所有符合条件数据的均值
# expanding().mean()实现分组内的累积均值，shift(1)排除当前行
data['lastavrage'] = data.groupby('ACOUNTNO')['avg_filtered'].expanding().mean().groupby('ACOUNTNO').shift(1)

# 处理"之前没有符合条件数据"的情况，把NaN替换为0
data['lastavrage'] = data['lastavrage'].fillna(0)

# 用向量化条件赋值生成label和problem列
mask = (data['average'] < data['lastavrage']) & (data['LASTREAD'] > 0)
data['label'] = np.where(mask, 'abnormal', data.get('label', ''))
data['problem'] = np.where(mask, 'error', data.get('problem', ''))

# 清理中间辅助列
data.drop('avg_filtered', axis=1, inplace=True)

关键逻辑说明： #

• 用groupby+expanding.mean实现分组内的滚动均值计算，这是Pandas底层优化过的操作，比Python循环快得多
• shift(1)确保只取当前行之前的历史数据，和原代码中data.head(i)的逻辑完全一致
• 用np.where做批量条件赋值，避免逐行判断的开销

注意事项： #

• 如果你的DataFrame索引不是连续整数，建议先执行data.reset_index(drop=True)再运行上述代码，确保shift操作的准确性
• 可以先取1000行小样本测试，验证结果和原循环代码一致后，再跑全量数据

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