Pandas行过滤与聚合问题：DataFrame计算出现KeyError报错

需求与问题背景 #

需要基于两个DataFrame（df1、df2）按以下规则生成统计结果：

• With E：统计df1中对应Acc下，indi_val末尾为'_E'的Amt总和；
• Without E：统计df1中对应Acc下，indi_val末尾不为'_E'的Amt总和；
• Normal：统计df1中对应Acc的所有Amt总和。

现有代码运行时抛出KeyError: 'Indi'错误，同时逻辑存在漏洞，需一并修复。

数据源 #

import pandas as pd

# df1数据
data1 = {
    'Acc': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 4],
    'indi_val': ['Val1', 'val2', 'Val_E', 'Val1_E', 'Val1', 'Val3', 'val2', 'val2_E', 'val22_E', 'val2_A', 'val2_V', 'Val_E', 'Val_A', 'Val', 'Val2', 'val7'],
    'Amt': [10, 20, 5, 5, 22, 38, 15, 25, 22, 23, 24, 56, 67, 45, 87, 88]
}
df1 = pd.DataFrame(data1)

# df2数据
data2 = {
    'Acc': [1, 1, 2, 2, 3, 4],
    'Indi': ['With E', 'Without E', 'With E', 'Without E', 'Normal', 'Normal']
}
df2 = pd.DataFrame(data2)

错误信息 #

get_loc
    raise KeyError(key)
KeyError: 'Indi'

错误原因分析 #

1. 列名大小写错误：原代码中引用df1的金额列时写的是df1row['amt']，但df1的列名实际是Amt（大写开头），导致KeyError；
2. 逻辑漏洞：Normal分支未过滤对应Acc，会累加df1所有行的Amt，不符合需求；
3. 低效写法：使用df1.apply循环append的方式性能极差，违背Pandas矢量化运算的设计思路。

修复后的解决方案 #

方案1：预处理统计（高效推荐） #

先提前计算每个Acc的三类统计值，再映射到df2，避免逐行循环：

# 预处理df1，计算每个Acc的三类统计值
with_e = df1[df1['indi_val'].str.endswith('_E')].groupby('Acc')['Amt'].sum().rename('With E')
total = df1.groupby('Acc')['Amt'].sum().rename('Total')
without_e = (total - with_e).rename('Without E')
normal = total.rename('Normal')

# 合并成统计结果表
stats = pd.concat([with_e, without_e, normal], axis=1).reset_index()

# 将统计结果映射到df2
def get_amt(row):
    return stats.loc[stats['Acc'] == row['Acc'], row['Indi']].values[0]

df2['Amt'] = df2.apply(get_amt, axis=1)

print(df2)

方案2：优化原函数逻辑 #

保留apply写法，但修复错误并优化过滤逻辑：

def get_indi(row):
    acc = row['Acc']
    indi_type = row['Indi']
    # 先过滤当前Acc的所有行
    filtered_df = df1[df1['Acc'] == acc]
    
    if indi_type == "With E":
        return filtered_df[filtered_df['indi_val'].str.endswith('_E')]['Amt'].sum()
    elif indi_type == "Without E":
        return filtered_df[~filtered_df['indi_val'].str.endswith('_E')]['Amt'].sum()
    elif indi_type == "Normal":
        return filtered_df['Amt'].sum()
    return 0

df2['Amt'] = df2.apply(get_indi, axis=1)

print(df2)

最终结果 #

运行后df2的输出为：

Acc        Indi  Amt
0    1      With E   10
1    1  Without E   90
2    2      With E   47
3    2  Without E   62
4    3      Normal  255
5    4      Normal   88

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