基于source优先级筛选含ID、year的Pandas DataFrame方法问询

我有一个包含ID、year、number、source四列的Pandas DataFrame，source1和source2可能存在ID与year重复的情况，source2内也可能存在同一ID和year的多行数据。示例数据如下：

import pandas as pd

data = {
    'ID': [1, 2, 1, 1, 1, 2, 2, 3],
    'year': [2021, 2021, 2021, 2021, 2022, 2022, 2022, 2021],
    'number': [400, 300, 450, 200, 500, 600, 100, 700],
    'source': ['source1', 'source1', 'source2', 'source2', 'source2', 'source2', 'source2', 'source2']
}

df = pd.DataFrame(data)

筛选规则 #

• 若同一ID和year存在source1数据，则保留source1行，删除对应ID和year的source2行；仅当同一ID和year无source1数据时，才保留source2行。
• 若使用source2数据，同一ID和year的所有行均需保留。

预期输出：原数据中索引2、3的行被删除（对应ID+year存在source1），其余source2行保留（对应ID+year无source1或year不同）。

尝试过的方法 #

1. 示例有效但大规模数据失效的代码：

source1Mask = df["source"] == "source1"
source2Mask = (df["source"] == "source2") & ~((df["ID"].isin(df[source1Mask]["ID"])) & (df["year"].isin(df[source1Mask]["year"])))
finalMask = source1Mask | source2Mask 
filtered_df = df[finalMask]

问题：isin会单独匹配所有出现过的ID和year，而非同时匹配ID+year组合，导致错误过滤掉ID相同但year不同的source2数据。

2. 有效但效率较低的代码：

grouped = df.groupby(['ID', 'year'])
df['source1count'] = grouped['source'].transform(lambda x: (x == 'source1').sum()).fillna(0)
df['source2count'] = grouped['source'].transform(lambda x: (x == 'source2').sum()).fillna(0)
filtered_df = df[(df["source"] == "source1") | ((df["source1count"] == 0) & (df["source2count"] >= 1))]

问题：两次transform+sum运算在大规模数据集上性能损耗明显。

高效解决方案 #

方案一：分组标记法（简洁高效）

利用groupby.transform结合any快速标记每个ID+year组是否存在source1，仅需一次分组操作，性能远优于求和统计：

# 按ID+year分组，标记每组是否存在source1
df['has_source1'] = df.groupby(['ID', 'year'])['source'].transform(lambda x: (x == 'source1').any())

# 执行筛选：保留source1，或无source1的source2数据
filtered_df = df[(df['source'] == 'source1') | ((df['source'] == 'source2') & ~df['has_source1'])]

# 清理临时列
filtered_df = filtered_df.drop('has_source1', axis=1)

方案二：索引匹配法（超大规模数据最优）

利用Pandas索引的哈希表快速查找特性，避免临时列，在超大规模数据集上性能更优：

# 提取所有存在source1的ID+year组合（去重）
source1_pairs = df[df['source'] == 'source1'][['ID', 'year']].drop_duplicates()

# 构建筛选条件：source1 或 source2且不在source1_pairs组合中
mask = (df['source'] == 'source1') | (
    (df['source'] == 'source2') &
    ~df.set_index(['ID', 'year']).index.isin(source1_pairs.set_index(['ID', 'year']).index)
)

filtered_df = df[mask]

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