Pandas如何高效计算df1每行对应df2最相似n行的列平均值

核心思路是抛弃Python层的逐行迭代，用numpy向量化广播实现全量相似度计算，同时用部分排序代替全排序取Top n，性能可提升数十到上百倍。

优化后代码 #

import numpy as np
import pandas as pd

# 示例数据初始化和原逻辑完全一致
np.random.seed(0)
cols = ['value1', 'value2', 'value3', 'value4', 'value5']
idx1 = [f'id_{str(n)}' for n in np.random.randint(250, size=3)]
df1 = pd.DataFrame(np.random.randint(1, 6, size=(3, 5)), index=idx1, columns=cols)
idx2 = [f'id_{str(n)}' for n in np.random.randint(250, size=10)]
df2 = pd.DataFrame(np.random.randint(1, 6, size=(10, 5)), index=idx2, columns=cols)

n = 3
# ------------------- 优化逻辑开始 -------------------
# 1. 转numpy数组用于广播计算
arr1 = df1[cols].values
arr2 = df2[cols].values
df2_val1 = df2['value1'].values
df2_val2 = df2['value2'].values
M = len(arr1)
N = len(arr2)
actual_n = min(n, N)  # 处理n大于df2行数的边界情况

# 2. 向量化计算相似度矩阵：形状为(M, N)，每个元素是df1第i行和df2第j行的相似度
sim_mat = (arr1[:, None] == arr2).sum(axis=2) * 0.2

# 3. 批量取每行Top n的索引，用argpartition做部分排序，比全排序效率高很多
# 加负号是因为argpartition默认取最小的k个，我们需要取最大的k个
top_n_indices = np.argpartition(-sim_mat, kth=actual_n, axis=1)[:, :actual_n]

# 4. 批量计算三个平均指标
avg_similarity = np.take_along_axis(sim_mat, top_n_indices, axis=1).mean(axis=1)
avg_val1 = df2_val1[top_n_indices].mean(axis=1)
avg_val2 = df2_val2[top_n_indices].mean(axis=1)

# 5. 结果写入df1
df1['avg_similarity'] = avg_similarity
df1['avg_val1'] = avg_val1
df1['avg_val2'] = avg_val2

优化点说明 #

• 完全替换iterrows循环：用numpy广播一次计算所有行对的相似度，所有计算都在C层执行，无Python层循环开销
• 部分排序替代全排序：np.argpartition不需要对df2所有行完全排序，仅保证前n个是最大值，数据量越大效率提升越明显
• 直接用numpy数组取值计算平均值，避免pandas行索引的额外开销

如果df1和df2行数特别大（比如超过10万行），内存放不下完整的(M,N)相似度矩阵，可以再做分块处理，把df1拆成小批次分别计算即可。

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