Pandas数值DataFrame每行与所有行高效比较的最优实现方案

方案选型结论 #

最优方案为NumPy向量化广播实现，性能远高于for循环、apply/applymap类方法：

• for循环、apply类方法：存在Python层循环开销，时间复杂度为O(n²*m)（n为行数、m为列数），数据量稍大就会严重卡顿
• 向量化实现：底层为C语言实现的运算逻辑，无Python层循环开销，执行速度最快、计算成本最低，相同数据量下性能比循环方案高10~100倍

完整实现代码 #

import pandas as pd
import numpy as np

# 示例数据
A = pd.DataFrame([[1,1,1,1,1],[1,1,1,1,1],[2,2,2,2,2], [1,10,1,1,1]])
thr = 0.3
B = A * thr

# 核心运算
A_arr = A.values
B_arr = B.values
n_rows = A.shape[0]

# 广播比较得到所有行对的匹配结果
mask = (A_arr[:, None] > B_arr).all(axis=-1)
# 排除同索引行（匹配示例输出需求，若不需要可删除该行）
mask[np.arange(n_rows), np.arange(n_rows)] = False

# 生成目标字典
result = {i: np.where(mask[i])[0].tolist() for i in range(n_rows)}

输出result即可得到你要求的示例结果：

{0: [1, 2], 1: [0, 2], 2: [0, 1], 3: [0, 1, 2]}

逻辑说明 #

1. 先将DataFrame转为NumPy数组，规避pandas索引带来的额外开销
2. 利用NumPy广播机制将两个二维数组扩展为三维数组，一次运算即可得到所有行对的逐元素比较结果
3. 对列维度执行all判断，筛选出所有元素都满足小于条件的行对
4. 若需要匹配示例输出，手动屏蔽同索引行的匹配结果
5. 最后根据布尔掩码提取符合要求的行索引，生成目标字典

注意事项 #

如果行数超过5万，三维数组会占用较多内存，可以考虑按行分块计算避免内存溢出，普通业务场景（行数<2万）直接使用上述代码即可。

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