Pandas DataFrame查找参考值上下最近邻值的实现方法

核心需求 #

• 以DataFrame中price列为参考值，从所有列名以gain开头的列中提取两个值：
  • closestdown：小于参考值的最接近值（最近下界）
  • closestup：大于参考值的最接近值（最近上界）
• 已知前提：所有gain开头的列已按行维度升序排列，每行的最近上界永远在最近下界的右侧相邻列，可利用该特性简化计算

示例数据 #

原始DataFrame结构：

SYMBOL        price   gainddS8  gainddS7_5   gainddS7  gainddS6_5  
102  1000SHIBUSDT     0.016049  -1.1       -0.1        1.5        2.12    
9         ADAUSDT     0.572700 -15.371514  -2.5       -1.0        2.497339   
24       ALGOUSDT     0.391300 -1.117796    0.5104497  14.091197   16.077897

期望输出结构（新增两列）：

SYMBOL        price  closestdown closestup
102  1000SHIBUSDT     0.016049 -0.1         1.5
9         ADAUSDT     0.572700 -1.0         2.497339
24       ALGOUSDT     0.391300  -1.117796   0.5104497

原有仅能查找全局最近值的代码：

df1 = df.filter(like='gain')
pos = df1.sub(df['price'], axis=0).abs().to_numpy().argmin(axis=1)
df['closestvalue'] = df1.to_numpy()[np.arange(len(df1)), pos]

实现代码 #

利用gain列行内升序的特性，直接定位每行第一个大于price的列位置，无需全量计算差值：

import numpy as np

# 提取所有gain开头的列转为numpy数组提升计算效率
gain_mat = df.filter(like='gain').to_numpy()
# price列转为列向量，方便广播比对
price_vec = df['price'].to_numpy().reshape(-1, 1)

# 找到每行第一个大于price值的列索引，即为最近上界的位置
up_idx = (gain_mat > price_vec).argmax(axis=1)
# 上界左侧相邻位置即为最近下界的位置
down_idx = up_idx - 1

# 批量赋值生成新列
df['closestdown'] = gain_mat[np.arange(len(df)), down_idx]
df['closestup'] = gain_mat[np.arange(len(df)), up_idx]

逻辑说明 #

• 因为gain列每行从左到右从小到大排列，第一个比price大的值就是大于参考值的最接近值，它左边紧邻的值就是小于参考值的最接近值，完全匹配需求
• 该实现比全列求绝对值差再找最小值的方案性能更高，数据量越大优势越明显
• 如果存在某行所有gain值都小于price、或都大于price的边界场景，可以在获取索引后加判断逻辑补充默认值，常规场景下可直接运行得到正确结果

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