如何根据条件从DataFrame中选取最近的左侧和右侧值？

可复现数据 #

import pandas as pd

df2 = pd.DataFrame( {'tof': {4875: 553.347572418771, 4876: 554.4911639867169, 4877: 556.2920840265602, 4878: 560.4059228421942, 4879: 560.7254631042018, 4880: 563.528024491697, 4881: 566.9995989172061, 4882: 570.7775523817776, 4883: 572.0685266789887}, 'E': {4875: 21.390025636898983, 4876: 21.301354836365054, 4877: 21.16283071077996, 4878: 20.85142476306085, 4879: 20.82752461774972, 4880: 20.61965532001513, 4881: 20.366451134218167, 4882: 20.096164210524464, 4883: 20.005036194577794}})

数据样例：

tof          E
4875  553.347572  21.390026
4876  554.491164  21.301355
4877  556.292084  21.162831
4878  560.405923  20.851425
4879  560.725463  20.827525
4880  563.528024  20.619655
4881  566.999599  20.366451
4882  570.777552  20.096164
4883  572.068527  20.005036

需求说明 #

根据给定的E_provided值，选取两个目标行：

• 最近左侧点：E值≤E_provided的所有点中，E值最大的那个
• 最近右侧点：E值≥E_provided的所有点中，E值最小的那个

边界情况处理：

• 当E_provided小于所有E值时，取E值最小的两个相邻点
• 当E_provided大于所有E值时，取E值最大的两个相邻点

比如：

• 当E_provided=20.83时，期望获取：

4879  560.725463  20.827525  # 左侧点
4878  560.405923  20.851425  # 右侧点

• 当E_provided=20.1时，期望获取：

4881  566.999599  20.366451
4882  570.777552  20.096164

原有方法的问题 #

之前尝试的代码通过绝对值差排序取前2个点，在边界场景下会返回同侧点，不符合需求：

df2.iloc[(df2['E']-E_provided).abs().argsort()[:2]]

解决方案 #

以下是满足需求的通用实现：

def get_neighbor_points(df, E_provided):
    # 获取左侧点：E<=E_provided中E最大的点
    left_point = df[df['E'] <= E_provided].nlargest(1, 'E')
    # 获取右侧点：E>=E_provided中E最小的点
    right_point = df[df['E'] >= E_provided].nsmallest(1, 'E')
    
    # 处理边界：E_provided小于所有E值，取最小的两个E点
    if left_point.empty:
        return df.nsmallest(2, 'E')
    # 处理边界：E_provided大于所有E值，取最大的两个E点
    if right_point.empty:
        return df.nlargest(2, 'E')
    # 正常情况合并左右两个点
    return pd.concat([left_point, right_point])

测试示例 #

# 测试E_provided=20.83
print(get_neighbor_points(df2, 20.83))
# 输出：
#              tof          E
# 4879  560.725463  20.827525
# 4878  560.405923  20.851425

# 测试E_provided=20.1
print(get_neighbor_points(df2, 20.1))
# 输出：
#              tof          E
# 4881  566.999599  20.366451
# 4882  570.777552  20.096164

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