Pandas按组筛选Source对应Name无交集的分组方法

问题背景 #

假设存在如下DataFrame：

Group  |  Source  |  Name
___________________________
  A    |    X     |  Jolly
  A    |    X     |  Stone
  A    |    X     |  Jolly
  A    |    Y     |  Sand
  B    |    X     |  Sand
  B    |    X     |  Stone
  B    |    Y     |  Stone
  C    |    X     |  Sand
  C    |    X     |  Stone

需求为筛选出同时满足以下两个条件的Group对应的所有行：

• 组内同时包含X、Y两类Source
• Source=X对应的Name集合，与Source=Y对应的Name集合完全无交集

以上述示例为例，符合要求的只有Group A：X对应的Name集合为{Jolly, Stone}，Y对应的Name集合为{Sand}，二者无公共值；Group B因为两类Source都存在Name=Stone被排除，Group C因为仅存在Source=X被排除。
期望输出结果如下：

Group  |  Source  |  Name
___________________________
  A    |    X     |  Jolly
  A    |    X     |  Stone
  A    |    X     |  Jolly
  A    |    Y     |  Sand

现有实现评估 #

你当前的实现代码如下：

def find_no_intersection(df):
    return (
        len(df[df.Source == 'X'].Name.values) > 0 and 
        len(df[df.Source == 'Y'].Name.values) > 0 and 
        (
            len(
                set(df[df.Source == 'X'].Name.values) & 
                set(df[df.Source == 'Y'].Name.values)
            ) == 0
        )
    )

df.groupby(['Group']).filter(find_no_intersection)

该实现逻辑完全正确，可以得到预期结果，但存在几处可优化的点：

• 同一分组内重复执行了4次df[df.Source == 'X']/df[df.Source == 'Y']的子集筛选，存在不必要的计算开销
• 计算交集后判断长度为0的写法，不如Python集合原生的isdisjoint()方法高效——后者只要检测到第一个公共元素就会立刻返回结果，不需要遍历完所有元素生成完整交集

优化方案 #

方案1：精简groupby逻辑（推荐，可读性高） #

优化自定义函数，减少重复筛选，替换更高效的集合判断方法：

def find_no_intersection(group_df):
    x_names = set(group_df.loc[group_df.Source == 'X', 'Name'])
    y_names = set(group_df.loc[group_df.Source == 'Y', 'Name'])
    # 两个Source都存在，且Name集合无交集
    return bool(x_names) and bool(y_names) and x_names.isdisjoint(y_names)

result = df.groupby('Group', group_keys=False).filter(find_no_intersection)

这个版本和原逻辑完全一致，代码更简洁，在中小数据集上性能足够，可读性最好。

方案2：向量化实现（适合超大数据集） #

如果数据量达到千万级，可以用pandas向量化操作替代逐组Python循环，进一步提升性能：

# 先对(Group, Source, Name)去重，减少后续计算量
unique_pairs = df.drop_duplicates(['Group', 'Source', 'Name'])
# 找出所有组内同时出现在X、Y两个Source下的Name，对应的Group就是不符合要求的
overlap_groups = unique_pairs.merge(
    unique_pairs,
    on=['Group', 'Name'],
    suffixes=('_x', '_y')
).query("Source_x == 'X' and Source_y == 'Y'")['Group'].unique()
# 先筛选出同时有X、Y的Group，再排除掉有交集的Group
valid_groups = df.groupby('Group')['Source'].agg(set) \
                 .loc[lambda s: s >= {'X', 'Y'}].index \
                 .difference(overlap_groups)

result = df[df['Group'].isin(valid_groups)]

这个方案全程用pandas内置的向量化方法，没有逐组的Python层循环，大数据量下速度优势明显。

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