如何移除Pandas多级分类DataFrame中无跨col2共同col3值的行

输入数据 #

import pandas as pd

df = pd.DataFrame({
    'col1': ['P', 'P', 'P', 'P', 'P', 'Q', 'Q', 'Q', 'Q', 'Q', 'Q'],
    'col2': ['X', 'X', 'Y', 'Y', 'Y', 'X', 'X', 'X', 'X', 'Y', 'Y'],
    'col3': ['B', 'D', 'A', 'B', 'D', 'A', 'B', 'C', 'D', 'A', 'C']
})

初始DataFrame展示：

col1 col2 col3
0     P    X    B
1     P    X    D
2     P    Y    A
3     P    Y    B
4     P    Y    D
5     Q    X    A
6     Q    X    B
7     Q    X    C
8     Q    X    D
9     Q    Y    A
10    Q    Y    C

需求说明 #

需要移除部分行，使得每个col1类别下，col3的取值必须在该col1对应的所有col2类别中都存在：

• 对于col1=P，col2包含X和Y，只有B、D在两个col2类别中都出现，保留对应行；
• 对于col1=Q，col2包含X和Y，只有A、C在两个col2类别中都出现，保留对应行。

期望输出 #

col1 col2 col3
0     P    X    B
1     P    X    D
3     P    Y    B
4     P    Y    D
5     Q    X    A
7     Q    X    C
9     Q    Y    A
10    Q    Y    C

已尝试的思路 #

你已经完成了两步关键分析：

1. 按col1+col2分组，查看每组的col3唯一值：

df.groupby(["col1", "col2"])["col3"].unique()

输出：

col1  col2
P     X             [B, D]
      Y          [A, B, D]
Q     X       [A, B, C, D]
      Y             [A, C]

2. 用pd.crosstab生成存在性真值表，按col1分组判断col3是否在所有col2中存在：

x = pd.crosstab([df.col1, df.col2], df.col3)
(x>0).groupby(level=0).agg(all)

输出：

col3      A      B      C      D
col1
P     False   True  False   True
Q      True  False   True  False

完整解决方案 #

方法一：基于交叉表真值表筛选 #

利用已生成的真值表，提取有效col1-col3配对，再和原DataFrame合并筛选：

# 生成存在性交叉表并计算每个col1分组下col3的有效性
x = pd.crosstab([df.col1, df.col2], df.col3)
valid_mask = (x > 0).groupby(level=0).agg(all)
# 转换为(col1, col3)有效配对格式
valid_pairs = valid_mask.stack().reset_index()
valid_pairs = valid_pairs[valid_pairs[0]][['col1', 'col3']]
# 筛选原DataFrame
result = df.merge(valid_pairs, on=['col1', 'col3'])

方法二：基于集合交集筛选 #

按col1分组，计算该组内所有col2对应的col3集合的交集，再筛选原DataFrame：

# 按col1分组，计算每组内所有col2对应的col3集合的交集
valid_col3 = df.groupby('col1')['col3'].apply(
    lambda g: set.intersection(*g.groupby(df.loc[g.index, 'col2']).unique().map(set))
)
# 转换为有效配对格式
valid_pairs = valid_col3.explode().reset_index()
# 筛选原DataFrame
result = df.merge(valid_pairs, on=['col1', 'col3'])

两种方法最终都会得到你期望的输出结果。

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