pandas调用idxmin()后如何提取对应位置带原始符号的值

使用idxmin()获取两列数据绝对值最小值对应的列索引，实现代码：

rep["OffsetFrom"] = rep[["OffsetDates", "OffsetDays"]].abs().idxmin(axis=1).dropna()

运行后得到的结果样例：

OffsetDates OffsetDays  OffsetFrom
0       0.0     0.0         OffsetDates
1       1.0     1.0         OffsetDates
2       4.0     -3.0        OffsetDays
3       4.0     -3.0        OffsetDays
4       6.0     -1.0        OffsetDays
... ... ... ...
1165    0.0     0.0         OffsetDates
1166    0.0     0.0         OffsetDates
1167    0.0     0.0         OffsetDates
1168    0.0     0.0         OffsetDates
1169    0.0     0.0         OffsetDates
1170 rows × 3 columns

需要新增一列存储绝对值最小位置对应的原始数值，必须保留数值原本的正负符号。

以下写法不符合需求，因为对绝对值取min会丢失原始符号：

rep[["OffsetFrom", "Offset"]] = rep[["OffsetDates", "OffsetDays"]].abs().agg(['idxmin','min'], axis=1)

目前通过逐行apply的方式实现了需求：

def get_abs_min_keep_sign(x):
    return min(x.min(), x.max(), key=abs)

rep["Offset"] = rep[["OffsetDates", "OffsetDays"]].apply(get_abs_min_keep_sign, axis=1)

想要找到更简洁、性能更好的实现方式替代逐行apply。

你已经算出了每一行绝对值最小的列名存在OffsetFrom列，不需要重复计算，直接用矢量化方式按列名取原始值即可，性能远高于逐行apply：

• 兼容旧版pandas（<1.2.0）的写法，直接用lookup行级查找：

rep["Offset"] = rep.lookup(row_labels=rep.index, col_labels=rep["OffsetFrom"])

• 新版pandas推荐写法，用numpy索引实现，无版本兼容问题，速度最快：

import numpy as np
target_cols = ["OffsetDates", "OffsetDays"]
# 一次性算出每一行绝对值最小的列位置索引
col_pos = rep[target_cols].abs().to_numpy().argmin(axis=1)
# 直接通过numpy位置索引取原始值，全程无Python层循环
rep["OffsetFrom"] = [target_cols[i] for i in col_pos]
rep["Offset"] = rep[target_cols].to_numpy()[np.arange(len(rep)), col_pos]

两种写法都能完整保留原始数值的正负符号，数据量越大性能优势越明显，比逐行apply的实现快数十倍。

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