基于Pandas：匹配DataFrame列值并获取对应行的上下2行数据

现有两个Pandas DataFrame，定义如下：

import pandas as pd

data1 = {
    'score': [1, 2],
    'seconds': [1140, 2100],
}

data2 = {
    'prize': [5.5, 14.5, 14.6, 21, 23, 24, 26, 38, 39, 40, 50],
    'seconds': [840, 1080, 1380, 1620, 1650, 1680, 1700, 1740, 2040, 2100, 2160],
}

df1 = pd.DataFrame.from_dict(data1)
df2 = pd.DataFrame.from_dict(data2)

df1输出：

score  seconds
0      1     1140
1      2     2100

df2输出：

prize  seconds
0     5.5      840
1    14.5     1080
2    14.6     1380
3    21.0     1620
4    23.0     1650
5    24.0     1680
6    26.0     1700
7    38.0     1740
8    39.0     2040
9    40.0     2100
10   50.0     2160

针对df1中seconds列的每个值，在df2已排序且唯一的seconds列中找到匹配或最接近的行，同时获取该行的上下各2行数据。期望结果如下：

prize  seconds
0     5.5      840
1    14.5     1080 # 最接近1140的匹配项
2    14.6     1380
3    21.0     1620
7    38.0     1740
8    39.0     2040
9    40.0     2100 # 完全匹配2100
10   50.0     2160

利用Pandas的searchsorted函数定位目标值的插入位置，确定最接近的行索引后扩展上下2行范围，最后去重合并结果：

import pandas as pd

# 定义原始数据
data1 = {
    'score': [1, 2],
    'seconds': [1140, 2100],
}

data2 = {
    'prize': [5.5, 14.5, 14.6, 21, 23, 24, 26, 38, 39, 40, 50],
    'seconds': [840, 1080, 1380, 1620, 1650, 1680, 1700, 1740, 2040, 2100, 2160],
}

df1 = pd.DataFrame.from_dict(data1)
df2 = pd.DataFrame.from_dict(data2)

# 提取用于查找的数组
targets = df1['seconds'].values
seconds_arr = df2['seconds'].values

# 找到每个目标值的插入位置
indices = pd.np.searchsorted(seconds_arr, targets, side='left')

# 确定每个目标对应的最接近行索引
closest_indices = []
for idx, target in zip(indices, targets):
    if idx == 0:
        closest_idx = 0
    elif idx == len(seconds_arr):
        closest_idx = len(seconds_arr) - 1
    else:
        left_diff = target - seconds_arr[idx-1]
        right_diff = seconds_arr[idx] - target
        closest_idx = idx - 1 if left_diff <= right_diff else idx
    closest_indices.append(closest_idx)

# 收集所有需要保留的索引（上下2行，自动去重）
keep_indices = set()
for idx in closest_indices:
    start = max(0, idx - 2)
    end = min(len(df2) - 1, idx + 2)
    keep_indices.update(range(start, end + 1))

# 按原顺序提取结果
result = df2.loc[sorted(keep_indices)]
print(result)

1. searchsorted快速定位：利用有序数组的特性，快速找到目标值的插入位置，缩小最接近值的查找范围。
2. 确定最接近索引：对每个插入位置，比较左右相邻值与目标值的差值，选择更接近的索引；边界情况直接取首尾。
3. 扩展行范围：对每个最接近索引，计算上下2行的有效索引范围（避免超出DataFrame边界），用集合存储去重。
4. 生成结果：按原始顺序提取目标行，得到最终需求的结果。

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