Pandas如何筛选col0、col1分组下col2同时包含-1和1的所有行

核心逻辑 #

按col0、col1分组后，仅保留同时包含col2=-1和col2=1的分组下所有行，其余分组整组删除。

完整实现代码 #

import pandas as pd

# 构造示例DataFrame，val列直接做除法运算得到对应数值
df = pd.DataFrame([
    ['a1', 'b1', -1, 0/1],
    ['a1', 'b1',  1, 1/1],
    ['a1', 'b2', -1, 2/1],
    ['a1', 'b2',  1, 1/2],
    ['a2', 'b1',  1, 1/3],
    ['a2', 'b2', -1, 2/3],
    ['a2', 'b2',  1, 4/1]
], columns=['col0', 'col1', 'col2', 'val'])

# 核心筛选逻辑：用groupby.filter直接保留符合条件的分组
result = df.groupby(['col0', 'col1']).filter(lambda group: {-1, 1}.issubset(group['col2'].unique()))

print(result)

输出结果 #

col0 col1  col2       val
0   a1   b1    -1  0.000000
1   a1   b1     1  1.000000
2   a1   b2    -1  2.000000
3   a1   b2     1  0.500000
5   a2   b2    -1  0.666667
6   a2   b2     1  4.000000

替代方案（避免多级索引问题） #

如果不想处理分组带来的多级索引，可以用分组标记+布尔索引的方式实现：

# 给每个分组标记是否同时包含-1和1
df['keep_flag'] = df.groupby(['col0', 'col1'])['col2'].transform(lambda x: {-1,1}.issubset(x.unique()))
# 筛选标记为True的行，删除辅助列即可得到最终结果
result = df[df['keep_flag']].drop('keep_flag', axis=1)

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