Python中匹配两个DataFrame字段的双层for循环如何优化提速？

优化方案说明 #

你原来的双重循环时间复杂度为O(M*N)，两个4万行的表运算量高达16亿次，Python原生循环执行效率低，跑5小时是正常现象。用Pandas向量化合并操作可以把时间复杂度降到接近线性，全程仅需几秒到几十秒就能完成运算，逻辑和你原代码完全对齐。

完整实现代码 #

import pandas as pd
import numpy as np

# 1. 生成匹配用的辅助列：对x、y做取整处理，和原逻辑一致
red_masked_tracks['x_round'] = red_masked_tracks['x'].round()
red_masked_tracks['y_round'] = red_masked_tracks['y'].round()
green_masked_tracks['x_round'] = green_masked_tracks['x'].round()
green_masked_tracks['y_round'] = green_masked_tracks['y'].round()

# 2. 对红表按匹配键去重，保留第一个出现的行，和原代码break的逻辑对齐
red_unique = red_masked_tracks.drop_duplicates(
    subset=['frame', 'x_round', 'y_round'], 
    keep='first'
)

# 3. 按三个字段做合并
# 如果需要保留绿表所有行、未匹配项留0，把how参数改成'left'即可
merged_df = pd.merge(
    green_masked_tracks,
    red_unique[['frame', 'x_round', 'y_round', 'x', 'y', 'particle', 'mass']],
    on=['frame', 'x_round', 'y_round'],
    how='inner',
    suffixes=('_green', '_red')
)

# 如果用了left合并，新增这行把空值填为0，和原代码初始化逻辑一致
# merged_df = merged_df.fillna(0)

# 4. 提取目标数组，和你原代码输出完全对应
intensity = merged_df[['mass_red', 'mass_green']].values
location = merged_df[['x', 'y']].values
frame = merged_df[['frame']].values
particle = merged_df[['particle']].values

注意事项 #

如果匹配结果和预期不一致，优先检查两个表的frame、x_round、y_round三个字段的数据类型是否一致，统一转成int类型后再合并即可避免匹配错误。

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