如何加速Python循环？Pandas大数据集统计逻辑优化求助

你的代码在小数据集上可以正常运行，但处理50万行的超大数据集时速度极慢，核心原因是Python循环中嵌套了每次迭代的切片求和操作——这相当于每次循环都要遍历一遍后半段数据，时间复杂度达到O(m*n)（m是B=1的行数，n是总行数），完全无法应对大数据量。

下面提供两种优化方案，均基于pandas/numpy的向量化运算（底层为C实现），速度可提升数个数量级：

方案一：全向量化运算（最快实现） #

利用numpy的广播特性，一次性计算所有B=1行的统计结果，彻底避免Python循环：

import pandas as pd
import numpy as np

def calc_optimized(df):
    # 复制原数据，避免修改输入的原始数据
    df = df.copy()
    # 筛选B=1的行，提取关键索引和对应值
    mask_b = df['B'] == 1
    target_indices = df[mask_b].index
    bv_values = df.loc[mask_b, 'B.v'].values
    a_values = df.loc[mask_b, 'A'].values
    
    a_array = df['A'].values
    total_rows = len(a_array)
    
    # 构造掩码矩阵：每个目标行仅关注自身及之后的行
    post_mask = np.arange(total_rows) >= target_indices[:, np.newaxis]
    
    # 计算每个目标行往后，A列小于B.v的数量，并批量赋值'in'列
    count_in = np.sum((a_array < bv_values[:, np.newaxis]) & post_mask, axis=1)
    df.loc[target_indices, 'in'] = np.where(count_in == 1, 1, np.where(count_in >= 2, 2, 0))
    
    # 同理计算A列大于当前A值的数量，批量赋值'out'列
    count_out = np.sum((a_array > a_values[:, np.newaxis]) & post_mask, axis=1)
    df.loc[target_indices, 'out'] = np.where(count_out == 1, 1, np.where(count_out >= 2, 2, 0))
    
    return df

# 测试用例
if __name__ == "__main__":
    df1 = pd.DataFrame({
        'index': [0,1,2,3,4],
        'Time': ['2022-01-01','2022-01-02','2022-01-03','2022-01-04','2022-01-05'],
        'A': [234,456,323,576,234],
        'B': [0,1,0,1,0],
        'B.v': [0,234,0,323,0],
        'in': [0,0,0,0,0],
        'out': [0,0,0,0,0]
    })
    print(calc_optimized(df1))

方案二：内存友好型（适合超大规模数据集） #

如果B=1的行数过多，方案一的广播矩阵会占用大量内存（比如10万*50万的矩阵），可以用反向遍历+二分查找优化，大幅降低内存占用：

import pandas as pd
import numpy as np

def calc_memory_efficient(df):
    df = df.copy()
    mask_b = df['B'] == 1
    target_indices = df[mask_b].index
    # 创建映射字典，快速查找目标行的对应值
    bv_map = df.loc[mask_b, 'B.v'].to_dict()
    a_map = df.loc[mask_b, 'A'].to_dict()
    
    a_array = df['A'].values
    sorted_post_a = []  # 维护从当前行往后的排序后的A值列表
    
    # 从后往前遍历，逐步构建排序后的列表
    for idx in reversed(range(len(a_array))):
        current_val = a_array[idx]
        # 将当前值插入到排序列表的正确位置
        insert_pos = np.searchsorted(sorted_post_a, current_val)
        sorted_post_a.insert(insert_pos, current_val)
        
        # 如果当前行是B=1的目标行，计算统计值并赋值
        if idx in target_indices:
            bv = bv_map[idx]
            a_val = a_map[idx]
            
            # 统计往后A中小于bv的数量
            count_in = np.searchsorted(sorted_post_a, bv)
            # 统计往后A中大于a_val的数量
            count_out = len(sorted_post_a) - np.searchsorted(sorted_post_a, a_val, side='right')
            
            df.loc[idx, 'in'] = 1 if count_in == 1 else 2 if count_in >= 2 else 0
            df.loc[idx, 'out'] = 1 if count_out == 1 else 2 if count_out >= 2 else 0
    
    return df

# 测试用例
if __name__ == "__main__":
    df1 = pd.DataFrame({
        'index': [0,1,2,3,4],
        'Time': ['2022-01-01','2022-01-02','2022-01-03','2022-01-04','2022-01-05'],
        'A': [234,456,323,576,234],
        'B': [0,1,0,1,0],
        'B.v': [0,234,0,323,0],
        'in': [0,0,0,0,0],
        'out': [0,0,0,0,0]
    })
    print(calc_memory_efficient(df1))

优化细节补充 #

1. 消除冗余操作：原代码反复调用columns.get_loc和iloc，改用loc直接按索引访问，减少不必要的计算开销。
2. 避免重复计算：原代码对同一个统计条件计算两次（比如sum(pointvalue > postrates)在两个if判断中重复执行），优化后仅计算一次。
3. 底层运算替代Python循环：Python循环的运行开销极高，用numpy的向量化运算或二分查找将循环转移到底层C代码执行，效率提升显著。

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