基于字符串列筛选分组DataFrame行及KeyError报错解决

首先，我完全理解你的需求：按Year分组后，保留两类行——要么Animal2非空，要么Animal2为空但Animal1存在于该分组所有非空Animal2的动物集合中。同时需要修复当分组无有效Animal2值时触发的KeyError。

核心解决方案 #

我们可以分两步处理：先为每个分组构建目标动物集合，再基于集合过滤分组内的行，同时处理边界情况避免报错。

步骤1：导入依赖并准备示例数据

import pandas as pd

# 第一个示例输入
df = pd.DataFrame({
    'Year': [2002,2002,2002,2002,2010,2010,2010,2010],
    'Animal1': ['Dog','Mouse','Lion','Duck','Dog','Cat','Lion','Mouse'],
    'Animal2': ['Mouse,Lion','','','','Cat','','','']
})

# 报错示例输入
df_error = pd.DataFrame({
    'Year': [2002,2002,2002,2010,2010],
    'Animal1': ['Dog','Mouse','Lion','Dog','Cat'],
    'Animal2': ['Mouse','','','','']
})

步骤2：为每个分组构建Animal2的动物集合

这个函数会提取分组内所有非空Animal2的值，拆分逗号分隔的动物名称，去重后转为集合：

def build_animal_set(group):
    animals = []
    # 遍历非空的Animal2值，拆分逗号
    for val in group['Animal2'].dropna().str.strip():
        if val:  # 跳过空字符串
            animals.extend(val.split(','))
    # 去重并清理空格（比如处理"Mouse,Lion"拆分后的空格）
    return set(animal.strip() for animal in animals)

# 按Year分组生成集合字典
animal_sets = df.groupby('Year').apply(build_animal_set).to_dict()

步骤3：定义分组过滤函数并应用

这里的关键是用get()方法处理无有效集合的分组，避免KeyError：

def filter_group(group):
    # 获取当前分组的目标集合，无则返回空集合
    target_animals = animal_sets.get(group.name, set())
    # 构建过滤掩码：
    mask = (
        # 保留Animal2非空的行
        group['Animal2'].str.strip() != ''
        # 或者Animal2为空但Animal1在目标集合中
        | ((group['Animal2'].str.strip() == '') & group['Animal1'].isin(target_animals))
    )
    return group[mask]

# 应用过滤，group_keys=False避免保留分组索引
filtered_df = df.groupby('Year', group_keys=False).apply(filter_group)
print(filtered_df)

运行第一个示例会得到你期望的输出：

Year Animal1    Animal2
0  2002    Dog   Mouse,Lion
1  2002  Mouse            
2  2002   Lion            
3  2010    Dog         Cat
4  2010    Cat            

修复报错的关键原因 #

你之前的代码触发KeyError: (2010, 'Dog')，是因为当某个分组（比如报错示例中的2010组）没有非空的Animal2值时，你的sets字典里不会有这个分组的键。直接用sets[g.name]会找不到键而报错。

改用animal_sets.get(group.name, set())后，当分组无有效集合时会返回空集合，此时Animal1不在空集合中，对应的行都会被过滤，正好符合需求——这类分组没有可参照的Animal2值，所以所有Animal2为空的行都应该被移除。

用报错示例测试：

animal_sets_error = df_error.groupby('Year').apply(build_animal_set).to_dict()
filtered_error_df = df_error.groupby('Year', group_keys=False).apply(filter_group)
print(filtered_error_df)

得到期望输出：

Year Animal1 Animal2
0  2002    Dog    Mouse
1  2002  Mouse         

额外说明 #

• 处理了Animal2中的空格问题（比如"Mouse, Lion"拆分后自动清理空格）
• 兼容了分组无任何非空Animal2的边界情况
• 保留了原始DataFrame的行顺序和结构

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