基于另一列值提取对应列数据的Pandas实现问题

基础场景：根据category匹配对应列值生成count列 #

给定包含category、foo_count、bar_count、bat_count的DataFrame，直接用lookup方法高效实现：

import pandas as pd

# 示例数据
df = pd.DataFrame({
    'category': ['foo', 'bar', 'bat'],
    'foo_count': [10, 15, 5],
    'bar_count': [20, 25, 15],
    'bat_count': [30, 35, 25]
})

# 生成count列
df['count'] = df.lookup(df.index, df['category'] + '_count')

如果需要更直观的逻辑（适合新手理解），用apply：

df['count'] = df.apply(lambda row: row[row['category'] + '_count'], axis=1)

实际业务场景：合并两种格式的total数据 #

已知category与对应列名的映射关系，分情况处理：

• 当category非空时，取映射列的数值
• 当category为空时，直接用total列的值

代码实现 #

假设映射关系为category_map = {'One': 'One Total', 'Two': 'Two Total'}，示例数据如下：

df = pd.DataFrame({
    'category': ['One', 'Two', None],
    'total': [None, None, 300],
    'One Total': [100, 150, None],
    'Two Total': [200, 250, None]
})

方法一：分步处理（清晰易懂）

# 根据category匹配对应列
df['matched_total'] = df['category'].map(category_map).apply(lambda col: df[col] if col in df.columns else None)
# 填充空值，合并为统一列
df['unified_total'] = df['matched_total'].fillna(df['total'])

方法二：一步到位（简洁高效）

df['unified_total'] = df.apply(
    lambda row: row[category_map[row['category']]] if pd.notna(row['category']) else row['total'],
    axis=1
)

简化版（如果列名是category + ' Total'规律）

如果映射列名都是category后加 Total的格式，无需单独定义映射：

df['unified_total'] = df.apply(
    lambda row: row[f"{row['category']} Total"] if pd.notna(row['category']) else row['total'],
    axis=1
)

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