Python中如何优化寻找包含最多数据点的双列最优区间组合代码？

核心优化思路 #

• 替换iterrows()循环：pandas循环遍历效率极低，改用向量化操作实现全量计算
• 摒弃append()：每次append()都会生成新DataFrame，改用一次性构造结果集
• 移除冗余操作：去掉不必要的copy()和apply()，用原生列级比较替代逐行判断

优化后的代码 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'col1': [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15],
    'col2': [1,1,1,1.2,1.2,1.3,1.5,2,2,2.3,2.5,2.8,3,3,3],
    'col3': [2,2.2,2.3,4.2,1.6,1.6,1.3,1.4,1.5,1.7,2.4,2.8,2.9,3,4]
})

# 1. 批量生成所有样本对应的区间边界
df_intervals = df.assign(
    min_c2=df['col2'] - 0.25,
    max_c2=df['col2'] + 0.25,
    min_c3=df['col3'] - 0.25,
    max_c3=df['col3'] + 0.25
)

# 2. 向量化计算每个区间包含的样本数（与原代码保持开区间判断逻辑）
c2_in_range = (df['col2'].values[:, None] > df_intervals['min_c2'].values) & \
              (df['col2'].values[:, None] < df_intervals['max_c2'].values)
c3_in_range = (df['col3'].values[:, None] > df_intervals['min_c3'].values) & \
              (df['col3'].values[:, None] < df_intervals['max_c3'].values)
sum_test = (c2_in_range & c3_in_range).sum(axis=0)

# 3. 构造最终结果DataFrame
df_combinations = df_intervals[['min_c2', 'max_c2', 'min_c3', 'max_c3']].assign(sum_test=sum_test)

# 可选：快速定位包含最多样本的区间组合
best_combination = df_combinations.loc[df_combinations['sum_test'].idxmax()]
print(best_combination)

代码说明 #

1. 区间生成：用assign()一行完成所有样本的区间边界计算，替代原循环中的重复赋值
2. 向量化匹配：借助numpy广播机制，一次性完成所有区间与所有样本的匹配判断，效率远高于逐行apply()
3. 结果构造：直接基于区间数据集添加统计列，避免循环中反复append()的性能损耗

关于现成函数 #

pandas没有专门针对这类「样本区间匹配计数」的现成函数，但通过组合numpy广播和pandas原生列操作，就能实现比原代码简洁数倍的逻辑。如果处理超大规模数据，可以考虑用scipy.spatial.KDTree做近邻搜索优化，但对于当前需求，上述向量化方案已经足够高效简洁。

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