Python Pandas：如何按ID生成记录最值列名的新列？

原始数据 #

ID   | COUNT_COL_A | COUNT_COL_B | SUM_COL_A | SUM_COL_B
-----|-------------|-------------|-----------|------------
111  | 10          | 10          | 320       | 120
222  | 15          | 80          | 500       | 500
333  | 0           | 0           | 110       | 350
444  | 20          | 5           | 0         | 0
555  | 0           | 0           | 0         | 0
666  | 10          | 20          | 60        | 50

需求说明 #

• 新增列TOP_COUNT_2：对每个ID，提取所有COUNT_前缀列中值最大的列名；若所有COUNT_列值相同，用逗号分隔所有列名；若所有COUNT_列值均为0，填入NaN。
• 新增列TOP_SUM_2：逻辑同TOP_COUNT_2，针对SUM_前缀列处理。

解决方案 #

可以通过封装通用函数来复用逻辑，避免重复代码：

import pandas as pd
import numpy as np

# 构造示例DataFrame
df = pd.DataFrame({
    "ID": [111, 222, 333, 444, 555, 666],
    "COUNT_COL_A": [10, 15, 0, 20, 0, 10],
    "COUNT_COL_B": [10, 80, 0, 5, 0, 20],
    "SUM_COL_A": [320, 500, 110, 0, 0, 60],
    "SUM_COL_B": [120, 500, 350, 0, 0, 50]
})

def extract_top_columns(df, prefix):
    # 筛选指定前缀的列
    target_cols = [col for col in df.columns if col.startswith(prefix)]
    if not target_cols:
        return pd.Series([np.nan]*len(df))
    
    # 计算每行最大值、判断全0情况
    row_max = df[target_cols].max(axis=1)
    all_zero = (df[target_cols] == 0).all(axis=1)
    
    # 收集每行中等于最大值的列名，用逗号拼接
    top_col_names = df[target_cols].apply(
        lambda row: ", ".join([col for col in target_cols if row[col] == row_max[row.name]]),
        axis=1
    )
    
    # 全0行替换为NaN
    top_col_names = top_col_names.mask(all_zero, np.nan)
    return top_col_names

# 生成目标列
df["TOP_COUNT_2"] = extract_top_columns(df, "COUNT_")
df["TOP_SUM_2"] = extract_top_columns(df, "SUM_")

# 输出结果
print(df.to_string(index=False))

运行结果 #

ID  COUNT_COL_A  COUNT_COL_B  SUM_COL_A  SUM_COL_B          TOP_COUNT_2          TOP_SUM_2
111           10           10        320        120  COUNT_COL_A, COUNT_COL_B          SUM_COL_A
222           15           80        500        500           COUNT_COL_B  SUM_COL_A, SUM_COL_B
333            0            0        110        350                   NaN          SUM_COL_B
444           20            5          0          0           COUNT_COL_A                   NaN
555            0            0          0          0                   NaN                   NaN
666           10           20         60         50           COUNT_COL_B          SUM_COL_A

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