pandas如何查找A列值最接近的行并计算距离匹配B列值

针对大体量DataFrame的场景，用pandas.merge_asof实现最近邻匹配是最优选择，时间复杂度为O(n log n)，远快于暴力循环匹配。注意merge_asof要求连接键必须提前排序，同时需要额外加唯一行标识避免匹配到当前行本身。

完整实现代码 #

import pandas as pd

# 加载示例数据
df = pd.DataFrame({'A' : [1, 5, 7, 2, 3, 4], 'B' : [5, 2, 7, 5, 1, 9]})

# 1. 加唯一行ID，用于后续排除匹配到自身的情况
df = df.reset_index(drop=False).rename(columns={'index': 'row_id'})

# 2. 构造匹配用的右表，提前按连接键A排序，重命名匹配列方便后续识别
right_df = df[['row_id', 'A', 'B']].sort_values('A').rename(columns={'A': 'match_A', 'B': 'match_B'})
# 左表也必须按连接键A排序，满足merge_asof的输入要求
left_df = df.sort_values('A')

# 3. 执行最近邻连接
match_res = pd.merge_asof(
    left_df,
    right_df,
    on='A',
    direction='nearest',
    suffixes=('_left', '_right')
)

# 4. 计算距离，剔除匹配到当前行的记录，对每行保留距离最小的匹配结果
match_res['distance'] = (match_res['A'] - match_res['match_A']).abs()
match_res = (
    match_res[match_res['row_id_left'] != match_res['row_id_right']]
    .sort_values(['row_id_left', 'distance'])
    .drop_duplicates('row_id_left', keep='first')
)

# 5. 把匹配结果合并回原表，恢复原表的行顺序
final_df = (
    df.merge(match_res[['row_id_left', 'distance', 'match_B']], left_on='row_id', right_on='row_id_left')
    .drop(columns=['row_id', 'row_id_left'])
)

运行结果 #

对应给出的示例，最终输出的final_df如下，完全符合预期：

A  B  distance  match_B
0  1  5         1        5
1  5  2         1        9
2  7  7         2        2
3  2  5         1        5
4  3  1         1        9
5  4  9         1        1

注意事项 #

• 如果遇到多个a'和当前a距离相等的平局情况，上述代码默认保留排序后最先出现的匹配项，你可以修改drop_duplicates的keep参数调整选择规则
• 代码通过唯一行ID而非A列值排除自身匹配，即使存在多行A列取值完全相同的情况，也能正确匹配到其他行的记录，此时距离为0
• 千万不要省略排序步骤，merge_asof不会自动对输入表排序，未排序的输入会返回错误的匹配结果
• 如果你允许匹配到当前行本身，直接删除剔除row_id_left != row_id_right的那行代码即可

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