如何基于另一DataFrame两列筛选并拆分pandas数据集为两个子集

嘿，这个需求我太熟悉了！要快速生成你要的两个子集，核心是避开低效的逐行判断，充分利用pandas的向量化操作和集合运算——这俩是处理这类行归属问题的神器。

先给你理清楚思路：dftrn是df里同时出现在df1和df2的行，dftst是df里既不在df1也不在df2的行。下面直接上最快的实现方式，分两种场景适配不同数据规模：

一、普通数据集（中小规模） #

先拿你给的示例数据来演示，先模拟出df1和df2（假设它们是df的子集）：

import pandas as pd

# 你的示例DataFrame
df = pd.DataFrame({ 
    'product': ['prod1', 'prod1', 'prod1', 'prod2', 'prod2', 'prod2'], 
    'date': ['2017-01-01', '2017-02-01', '2017-03-01', '2017-02-01', '2017-03-01', '2017-04-01']
})

# 模拟两个子集df1和df2（你可以替换成自己的实际子集）
df1 = df[df['product'] == 'prod1']
df2 = df[df['date'].isin(['2017-02-01', '2017-03-01'])]

1. 生成dftrn（同时在df1和df2的行） #

用集合的交集运算快速定位共同行，再结合布尔索引筛选：

# 把df1和df2的行转成无索引元组，转成集合求交集
common_rows = set(df1.itertuples(index=False, name=None)) & set(df2.itertuples(index=False, name=None))

# 筛选df中属于交集的行
dftrn = df[df.itertuples(index=False, name=None).isin(common_rows)]

itertuples(index=False)能快速把每行转成不含索引的元组，集合的交集运算速度极快，再搭配isin的向量化判断，比merge或者逐行比对高效太多。

2. 生成dftst（既不在df1也不在df2的行） #

同样用集合运算，先合并df1和df2的所有行，再筛选df中不在这个集合里的行：

# 合并df1和df2的所有行到一个集合
combined_rows = set(df1.itertuples(index=False, name=None)) | set(df2.itertuples(index=False, name=None))

# 取反筛选不在合并集合中的行
dftst = df[~df.itertuples(index=False, name=None).isin(combined_rows)]

这里的~是布尔取反操作，正好对应“既不在df1也不在df2”的条件。

二、超大数据集（百万级以上） #

如果你的DataFrame数据量特别大，上面的方法虽然好用，但还有更极致的优化——利用复合索引的哈希查找特性，速度会再上一个台阶：

# 给df、df1、df2设置由所有列组成的复合索引
df_idx = df.set_index(['product', 'date'])
df1_idx = df1.set_index(['product', 'date'])
df2_idx = df2.set_index(['product', 'date'])

# 生成dftrn：取两个索引的交集，再恢复成原结构
common_index = df1_idx.index.intersection(df2_idx.index)
dftrn = df_idx.loc[common_index].reset_index()

# 生成dftst：取两个索引的并集，筛选不在并集中的行
combined_index = df1_idx.index.union(df2_idx.index)
dftst = df_idx.loc[~df_idx.index.isin(combined_index)].reset_index()

pandas的索引是基于哈希表实现的，交集、并集以及isin判断的效率都远高于行元组的集合操作，处理大数据集时优势非常明显。

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