基于Pandas高效查找DataFrame分组内对应前序数值的行索引

问题场景 #

现有两个结构一致的有序Pandas DataFrame：

• keys：包含分组列（group1、group2…groupN）和number列
• fill_in：同结构的有序DataFrame

需要为fill_in的每一行，在keys中找到满足以下条件的行索引：

1. 所有分组列的值完全匹配
2. number值小于当前行的number，且是该分组内满足条件的最大值

如果分组不存在，或当前行的number小于分组内所有number值，返回np.nan。

示例数据：

import pandas as pd
import numpy as np

keys = pd.DataFrame({'group1':[1, 1, 1, 1, 2, 2],
                     'group2':[5, 5, 5, 7, 9, 9],
                     'number': [19,35,61,5, 105,300]})
fill_in = pd.DataFrame({'group1':[1, 1, 2, 5],
                     'group2':[5, 5, 9, 9],
                     'number': [0,43.2,900.3,14]})
# 预期输出：[np.nan, 1, 5, np.nan]

解决方案 #

利用Pandas的merge_asof函数实现向量式操作，彻底替代循环，大幅提升性能：

步骤1：预处理数据 #

为两个DataFrame保留原始索引，并按分组列+number升序排序（merge_asof要求匹配列有序）：

# 为keys添加原始索引，并重命名便于区分
keys_with_idx = keys.reset_index().rename(columns={'index': 'keys_idx'})
keys_sorted = keys_with_idx.sort_values(by=['group1', 'group2', 'number'])

# 为fill_in添加原始索引，按相同规则排序
fill_in_with_idx = fill_in.reset_index().rename(columns={'index': 'fill_in_idx'})
fill_in_sorted = fill_in_with_idx.sort_values(by=['group1', 'group2', 'number'])

步骤2：分组匹配最大符合项 #

使用merge_asof在同分组内，查找小于当前number的最大匹配项：

merged = pd.merge_asof(
    fill_in_sorted,
    keys_sorted,
    on='number',               # 按number列匹配
    by=['group1', 'group2'],   # 限定仅同分组内匹配
    direction='backward',      # 查找小于当前number的最大项
    allow_exact_matches=False  # 直接排除number相等的情况
)

步骤3：恢复原始顺序并生成结果 #

按fill_in的原始索引排序，提取匹配到的keys索引，缺失值填充为np.nan：

result = merged.sort_values('fill_in_idx')['keys_idx'].replace({pd.NA: np.nan}).tolist()
print(result)  # 输出：[nan, 1.0, 5.0, nan]

关键说明 #

• merge_asof是Pandas专门用于有序数据匹配的函数，内部通过向量式操作实现，性能远优于Python循环
• direction='backward'确保找到当前行之前的最大匹配项，allow_exact_matches=False直接过滤掉number相等的情况，完全符合需求
• 预处理时的排序是merge_asof的必要条件，保证匹配逻辑的正确性

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