使用Pandas分组筛选同时满足多列至少一个True的id_p（求一行式解法）

原始数据 #

给定的pandas DataFrame代码如下：

import pandas as pd
foo = pd.DataFrame({'id_p': [1,1,2,2,3,3,3,4,4], 
              'id_d_b': [True, True, False, True, True, True,False,False,False], 
              'id_d_i': [False, False, True, False,False,False,True,True,True]})

数据展示：

需求 #

筛选出至少在id_d_b列有一个True，并且至少在id_d_i列有一个True的id_p。

现有实现 #

你已经实现了正确的逻辑，代码如下：

foo['id_d_b'] = foo['id_d_b'].astype(int)
foo['id_d_i'] = foo['id_d_i'].astype(int)

foo['has_id_d_b'] = foo.groupby('id_p')['id_d_b'].transform('max')
foo['has_id_d_i'] = foo.groupby('id_p')['id_d_i'].transform('max')

foo['result'] = foo['has_id_d_b'] + foo['has_id_d_i'] # 若值>1则该id_p满足条件
foo['result'] = foo.eval('result > 1')

执行后结果：

一行式解决方案 #

这里提供两种简洁的一行式写法，无需转换布尔值为整数：

写法1：利用transform与布尔运算 #

直接对两列分组后判断每组是否存在至少一个True，再做逻辑与：

foo['result'] = foo.groupby('id_p')['id_d_b'].transform('any') & foo.groupby('id_p')['id_d_i'].transform('any')

写法2：先筛选符合条件的id_p再标记 #

先通过filter提取满足条件的组，再用isin标记原表行：

foo['result'] = foo['id_p'].isin(foo.groupby('id_p').filter(lambda g: g['id_d_b'].any() and g['id_d_i'].any())['id_p'].unique())

两种写法都能得到和你现有实现完全一致的结果，且不需要修改原始数据类型。

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