优化Pandas数据集比对运行时：跨年度白宫访客交集查询提速

嘿，我来帮你搞定这个慢代码的问题！处理这种跨年份的访客交集，原函数速度慢大概率是因为用了低效的匹配逻辑（比如嵌套循环、重复全表扫描或者未优化的合并操作）。咱们可以用集合哈希或者pandas内置的分组聚合来大幅提速，下面给你两个靠谱的实现思路：

方案一：集合求交集（最快的轻量方案） #

集合的哈希查找是O(1)级别的，交集操作效率极高，适合只需要提取姓名的场景：

import pandas as pd

def task3_optimized():
    # 第一步：从TOA字段提取到访年份（先转成日期格式）
    data["VISIT_YEAR"] = pd.to_datetime(data["TOA"], errors="coerce").dt.year
    
    # 只保留2009和2010年的数据，减少后续处理量
    filtered_data = data[data["VISIT_YEAR"].isin([2009, 2010])]
    
    # 生成两年的访客姓名集合（用元组组合姓和名，避免同名不同人）
    visitors_2009 = set(zip(filtered_data[filtered_data["VISIT_YEAR"] == 2009]["NAMELAST"],
                            filtered_data[filtered_data["VISIT_YEAR"] == 2009]["NAMEFIRST"]))
    visitors_2010 = set(zip(filtered_data[filtered_data["VISIT_YEAR"] == 2010]["NAMELAST"],
                            filtered_data[filtered_data["VISIT_YEAR"] == 2010]["NAMEFIRST"]))
    
    # 求交集就是同时两年到访的访客
    common_visitors = visitors_2009.intersection(visitors_2010)
    
    # 转成DataFrame格式返回（如果需要）
    return pd.DataFrame(common_visitors, columns=["NAMELAST", "NAMEFIRST"])

为什么这个方案快？ #

集合的底层是哈希表，查找和交集操作的时间复杂度接近O(n)，比pandas的merge或者手动循环匹配快好几倍，尤其是数据量很大的时候优势更明显。

方案二：分组聚合（适合保留更多访客信息） #

如果之后还需要用到这些访客的其他数据（比如到访时间），可以用分组聚合的方式，一次性筛选出符合条件的访客：

import pandas as pd

def task3_groupby_optimized():
    data["VISIT_YEAR"] = pd.to_datetime(data["TOA"], errors="coerce").dt.year
    
    # 按姓名分组，收集每个访客的所有到访年份
    visitor_year_groups = data.groupby(["NAMELAST", "NAMEFIRST"])["VISIT_YEAR"].agg(lambda x: set(x.dropna()))
    
    # 筛选出年份集合同时包含2009和2010的访客
    common_visitors = visitor_year_groups[visitor_year_groups.apply(lambda years: 2009 in years and 2010 in years)]
    
    # 重置索引转成标准DataFrame
    return common_visitors.reset_index()

额外注意事项 #

• 如果TOA字段有异常日期值，记得用errors="coerce"把无效值转成NaN，避免后续报错。
• 用NAMELAST+NAMEFIRST作为唯一标识可能存在同名不同人的情况，如果数据集有其他标识字段（比如中间名、证件号），可以加入分组/集合的组合键来提升准确性。
• 处理超大数据集时，可以先把数据按年份拆分保存成小文件，再分别加载处理，减少内存占用。

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