如果包含特定值则筛选分组

Question

假设我有一个数据表，里面有成百上千个不同的名字：

df = pd.DataFrame({"NAME":["A", "A", "A" ,"A", 
                       "B", "B", "B", 
                       "C", "C", "C", "C",
                       "D", "D",],
               "CLASS":["LI","BO","EQ","AI",
                           "LI", "EQ", "AI", 
                           "LI","BO","EQ","AI",
                           "EQ", "AI", 
                          ]})

我想用Pandas的groupby功能，筛选出那些同时包含“EQ”和“AI”，但不包含“BO”的名字。理想情况下，我应该得到下面这样的结果：

NAME CLASS
B    LI
B    EQ
B    AI
D    EQ
D    AI

有没有什么建议？

Answer 1

使用 groupby.transform 结合集合操作：

out = df[df.groupby('NAME')['CLASS']
           .transform(lambda x: ((S:=set(x))>={'EQ', 'AI'}) and ('BO' not in S))]
# or
# out = df[df.groupby('NAME')['CLASS']
#            .transform(lambda x: ((S:=set(x))>={'EQ', 'AI'})
#                       and not S.intersection({'BO'}))]

另一种方法是使用 groupby.agg 和 isin：

g = df.groupby('NAME')['CLASS'].agg(set)
out = df[df['NAME'].isin(g.index[(g >= {'EQ', 'AI'}) & ~(g >= {'BO'})])]

输出结果：

   NAME CLASS
4     B    LI
5     B    EQ
6     B    AI
11    D    EQ
12    D    AI

中间结果

这里多加了一个组以便完整性

# transform approach
   NAME CLASS  g >= {'EQ', 'AI'}  ~(g >= {'BO'})      &
0     A    LI               True           False  False
1     A    BO               True           False  False
2     A    EQ               True           False  False
3     A    AI               True           False  False
4     B    LI               True            True   True
5     B    EQ               True            True   True
6     B    AI               True            True   True
7     C    LI               True           False  False
8     C    BO               True           False  False
9     C    EQ               True           False  False
10    C    AI               True           False  False
11    D    EQ               True            True   True
12    D    AI               True            True   True
13    E    AI              False            True  False

# agg approach
                     g  (g >= {'EQ', 'AI'})  ~(g >= {'BO'})      &
NAME                                                              
A     {LI, EQ, BO, AI}                 True           False  False
B         {LI, EQ, AI}                 True            True   True
C     {LI, EQ, BO, AI}                 True           False  False
D             {EQ, AI}                 True            True   True
E                 {AI}                False            True  False