Pandas DataFrame 如何实现行级A/B/C最大值对应列留值其余返回NaN

核心逻辑 #

• 先提取需要比较的A/B/C三列，计算每行的最大值
• 生成布尔掩码标记每个位置是否为该行最大值
• 非最大值位置替换为NaN，重命名为带-hat后缀的列后合并到原DataFrame即可

完整可运行代码 #

首先构造和你给出的样例一致的输入DataFrame：

import pandas as pd
import numpy as np

# 构造示例输入
data = {
    'A': [3.57, 4.98, 3.79, 3.95, -3.10],
    'B': [5.29, 9.64, 5.25, 5.65, -7.20],
    'C': [6.23, 7.58, 6.26, 6.61, -8.16]
}
df = pd.DataFrame(data, index=pd.to_datetime(['2020-01-05', '2020-01-04', '2020-01-03', '2020-01-02', '2020-01-01']))
df.index.name = 'Date'

核心实现代码（采用向量化操作，数据量大时效率远高于逐行遍历）：

# 定义原始列和对应的hat列名
src_cols = ['A', 'B', 'C']
hat_cols = [f'{col}-hat' for col in src_cols]

# 计算每行最大值
row_max = df[src_cols].max(axis=1)
# 生成掩码：标记每个位置是否为该行最大值
mask = df[src_cols].eq(row_max, axis='index')
# 非最大值位置设为NaN，重命名后赋值给新列
df[hat_cols] = df[src_cols].where(mask, np.nan).set_axis(hat_cols, axis=1)

输出验证 #

运行print(df)即可得到和期望完全一致的结果：

A     B     C  A-hat  B-hat  C-hat
Date                                              
2020-01-05  3.57  5.29  6.23    NaN    NaN  6.23
2020-01-04  4.98  9.64  7.58    NaN   9.64    NaN
2020-01-03  3.79  5.25  6.26    NaN    NaN  6.26
2020-01-02  3.95  5.65  6.61    NaN    NaN  6.61
2020-01-01 -3.10 -7.20 -8.16  -3.10    NaN    NaN

补充说明 #

• 如果存在某行多列并列最大值的场景，默认实现会让多个对应hat列都返回该最大值；如果需要强制每行仅保留1个有效值（哪怕并列也只取第一个出现的最大值），可以在生成掩码后加一步去重：

  # 仅保留同一行第一个出现的最大值标记
  mask = mask.mask(mask.cumsum(axis=1).gt(1), False)
  

• 如果使用的pandas版本低于1.3.0，可将重名列的代码替换为兼容性更好的写法：

  df[hat_cols] = df[src_cols].where(mask, np.nan).rename(columns=dict(zip(src_cols, hat_cols)))
  

• 逐行apply的写法虽然逻辑直观，但数据量超过万行时效率极低，不推荐使用。

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