如何向量化循环中的‘if not:’语句以优化Pandas代码效率？

我正在优化现有代码以提升效率，目前用iterrows()遍历DataFrame，每个文件都要执行这个操作，导致程序耗时太长。这个循环遍历DataFrame的每一行（仅含2列），如果‘3’列的值在另一个DataFrame中找不到就跳过，否则执行简单计算。

以下是待优化的代码，它返回两个列表（BP_Shear和BP_Dis）用于后续绘图：

BP_Shear = []
BP_Dis = []
for index, Total_row in Total_Data.iterrows():
    if not IT_Data[IT_Data['3'] == Total_row['3']].empty:
        BP_Shear.append(Total_row['3'])

        row_number = IT_Data[IT_Data['3'] == Total_row['3']].index

        BP_Dis.append(Total_row['2'] - mean(IT_Data.iloc[row_number]["2"]))

我试过用列表推导，但耗时反而增加了20秒，现在想实现向量化操作但没成功。

核心逻辑 #

遍历Total_Data的每一行（代码中为Total_row），先检查该行‘3’列的值是否存在于IT_Data中，不存在则跳过；若存在，将Total_row['3']加入BP_Shear，找到IT_Data中‘3’列值匹配的所有行索引，计算这些行‘2’列的均值，用Total_row['2']减去该均值，将结果加入BP_Dis。

示例数据 #

IT_Data #

Total_Data #

优化方案：向量化操作实现 #

步骤1：预计算分组均值 #

先对IT_Data按‘3’列分组，计算每组‘2’列的均值，避免循环中重复查询计算：

it_mean = IT_Data.groupby('3')['2'].mean()

步骤2：筛选有效行 #

用isin()快速筛选Total_Data中‘3’列值存在于IT_Data中的行：

filtered_total = Total_Data[Total_Data['3'].isin(it_mean.index)]

步骤3：批量计算结果 #

通过映射均值到筛选后的行，直接批量计算差值：

BP_Shear = filtered_total['3'].tolist()
BP_Dis = (filtered_total['2'] - filtered_total['3'].map(it_mean)).tolist()

完整优化代码 #

# 预计算IT_Data中‘3’列分组的‘2’列均值
it_mean = IT_Data.groupby('3')['2'].mean()
# 筛选Total_Data中符合条件的行
filtered_total = Total_Data[Total_Data['3'].isin(it_mean.index)]
# 生成结果列表
BP_Shear = filtered_total['3'].tolist()
BP_Dis = (filtered_total['2'] - filtered_total['3'].map(it_mean)).tolist()

效率说明 #

• 彻底避免了iterrows()的逐行开销，所有操作都是Pandas原生向量化运算，速度会大幅提升。
• 预计算分组均值只执行一次，解决了原代码中循环内重复布尔索引查询的性能瓶颈（原代码每次循环做两次IT_Data[IT_Data['3'] == ...]查询）。

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