优化低效Python循环：DataFrame嵌套For循环加速方案咨询

我没有正规软件培训背景，正在做个人项目，程序功能正常，但一段代码效率极低，整体执行要2天才能完成。

场景说明 #

• 长字符串列表，元素仅为'A'、'B'、'C'、'D'之一
• 两个DataFrame：
  • df_1：存储起始与结束区间值
  • df_2：统计列表中各字符串在对应区间内的出现次数

当前低效代码 #

import pandas as pd

# Create the test list
test_list = []
for i in range(0, 250):
    test_list.append('A')
    test_list.append('B')
    test_list.append('C')
    test_list.append('D')

# Create the test dataframes
data_1 = []
columns_1 = ['start', 'end']
for i in range(0, 10):
    data_1.append([i, i+10])    
df_1 = pd.DataFrame(data_1, columns=columns_1)

data_2 = []
columns_2 = ['A', 'B', 'C', 'D']
for i in range(0, 20):
    data_2.append([0, 0, 0, 0])
df_2 = pd.DataFrame(data_2, columns=columns_2)

# The checking loop
for index, row in df_2.iterrows():
    for x in range(0, len(test_list)):
        for ind, r in df_1.iterrows():
            if r['start'] <= x <= r['end']:
                if test_list[x] == 'A':
                    df_2.at[index, 'A'] +=1
                if test_list[x] == 'B':
                    df_2.at[index, 'B'] +=1
                if test_list[x] == 'C':
                    df_2.at[index, 'C'] +=1
                if test_list[x] == 'D':
                    df_2.at[index, 'D'] +=1

问题根源 #

三重嵌套循环（遍历df_2 + 遍历列表索引 + 遍历df_1区间），时间复杂度为O(N×M×K)（N=df_2行数，M=列表长度，K=df_1区间数），再加上df_2.at逐元素修改的开销，数据量大时直接导致性能爆炸。

优化方案 #

1. 用向量化操作替代嵌套循环

pandas/numpy的向量化操作是性能提升的核心，避免逐元素遍历：

方法一：区间匹配+分组统计（适合大列表）

先把列表转成带索引的DataFrame，再通过区间匹配关联到df_1的区间，最后分组统计：

import pandas as pd

# 预处理列表为带索引的DataFrame
test_df = pd.DataFrame({
    'index': range(len(test_list)),
    'char': test_list
})

# 给df_1添加区间对应的df_2行索引
df_1['df2_index'] = df_1.index

# 用merge_asof匹配每个索引所属的区间（需确保df_1按start排序）
df_1_sorted = df_1.sort_values('start')
merged = pd.merge_asof(
    test_df.sort_values('index'),
    df_1_sorted[['start', 'end', 'df2_index']],
    left_on='index',
    right_on='start',
    direction='backward'
)

# 过滤掉不在区间end范围内的记录
merged = merged[merged['index'] <= merged['end']]

# 分组统计每个区间的字符次数
counts = merged.groupby(['df2_index', 'char']).size().unstack(fill_value=0)

# 将统计结果批量写入df_2
df_2.loc[counts.index, counts.columns] = counts

方法二：numpy广播矩阵运算（适合区间数量少的场景）

利用numpy广播一次性判断所有索引的区间归属，再通过矩阵乘法计算统计结果：

import pandas as pd
import numpy as np

# 转换为numpy数组提升运算速度
indices = np.arange(len(test_list))
starts = df_1['start'].values
ends = df_1['end'].values

# 广播判断每个索引是否属于每个区间（形状：列表长度 × 区间数）
in_interval = (indices[:, None] >= starts) & (indices[:, None] <= ends)

# 将字符串列表转为one-hot编码（形状：列表长度 × 4）
chars_one_hot = pd.get_dummies(np.array(test_list)).values

# 矩阵乘法计算每个区间的字符次数（形状：区间数 × 4）
result = in_interval.T @ chars_one_hot

# 批量赋值到df_2
df_2.iloc[:len(result), :] = result

2. 避免低效操作

• 不要用df.iterrows()：它会把每行转成Series，效率极低，优先用向量化或df.apply()（仅万不得已时用）
• 不要用df.at逐元素修改：批量赋值（如df.loc[]、df.iloc[]）的速度是逐元素修改的百倍以上

3. 多进程优化（超大数据量场景）

如果单进程仍无法满足需求，可以拆分任务到多进程处理：

import pandas as pd
import numpy as np
from concurrent.futures import ProcessPoolExecutor

def process_single_interval(start, end, indices, chars_one_hot):
    # 判断当前区间内的所有索引
    in_range = (indices >= start) & (indices <= end)
    # 返回该区间的字符统计结果
    return in_range @ chars_one_hot

# 预处理数据为numpy数组（方便进程间传递）
indices = np.arange(len(test_list))
chars_one_hot = pd.get_dummies(np.array(test_list)).values

# 生成任务列表：每个区间对应一个任务
tasks = [(row['start'], row['end']) for _, row in df_1.iterrows()]

# 多进程执行任务
with ProcessPoolExecutor() as executor:
    results = list(executor.map(lambda x: process_single_interval(x[0], x[1], indices, chars_one_hot), tasks))

# 将结果写入df_2
df_2.iloc[:len(results), :] = np.array(results)

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