如何用groupby筛选满足双docn条件的Pandas DataFrame

给定如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame(data=[['840', '007658', 'EP010A', 'wrwrwr'],
 ['841', '007658', 'EP019410A', 'wwtert'],
 ['842', '007658', 'EP0129W', 'erterte'],
 ['843', '007658', 'EP0629W', 'ertetet'],
 ['992', '007675', 'EP0275A', 'ete'],
 ['993', '007675', 'EP0375A', 'ertre'],
 ['994', '007675', 'EP02091', 'ert'],
 ['117', '007690', 'EP02212A', 'sf'],
 ['118', '007690', 'EP00212A', 'sf'],
 ['118', '007690', 'EP02281W', 'dg'],
 ['118', '007690', 'EP07281W', 'dg'],
 ['118', '007690', 'EP0281W', 'sdf'],
 ['118', '007690', 'EP02281W', 'ert'],
 ['118', '007690', 'EP0612A', 'dfg'],
 ['11', '0076', 'US0612A', 'dfg'],
 ['12', '0076', 'CA0612A', 'dfg']], columns = ['i', 'fn','docn','ad']) 

需求：生成新的DataFrame，包含所有满足以下条件的行：同一fn编号下，至少存在一个以EP开头且不以W结尾的docn，同时至少存在一个以EP开头且以W结尾的docn。

举例说明：

• fn=007658包含EP010A（EP开头非W结尾）和EP0129W（EP开头W结尾），符合条件；
• fn=007675没有EP开头且W结尾的docn，不符合；
• fn=0076没有EP开头的docn，不符合。

用户尝试的代码无法实现需求：

for i,group in enumerate(df2.groupby("fn")):
    display(group)
    has_EP  = group[1]['docn'].apply(lambda docns: any(docn.startswith('EP') and not docn.endswith('W') for docn in docns))
    has_EP_W= group[1]['docn'].apply(lambda docns: any(docn.startswith('EP') and docn.endswith('W') for docn in docns))
    total_condition =has_EP & has_EP_W
    
    display(group[1]['docn'].loc[total_condition])
    
    if i==3:
        break

问题分析 #

你之前的代码错误在于：对分组后的docn列使用apply时，每个docns其实是单个字符串（而非分组内的所有docn集合），所以any()判断的是单个字符串是否满足条件，而非整个分组是否存在满足条件的项，导致逻辑完全错误。

方法一：使用groupby.filter()（最简洁） #

filter()方法可以直接对分组进行条件判断，保留满足条件的整个分组：

# 定义分组筛选条件
def filter_fn(group):
    # 检查分组内是否存在EP开头且非W结尾的docn
    has_ep_non_w = group['docn'].str.startswith('EP') & ~group['docn'].str.endswith('W')
    # 检查分组内是否存在EP开头且W结尾的docn
    has_ep_w = group['docn'].str.startswith('EP') & group['docn'].str.endswith('W')
    # 两个条件同时满足
    return has_ep_non_w.any() and has_ep_w.any()

# 应用筛选，得到结果
result_df = df.groupby('fn').filter(filter_fn)

方法二：先计算分组标志，再合并筛选 #

先对每个fn计算是否符合条件，再通过isin筛选原DataFrame：

# 计算每个fn的两个条件标志
fn_flags = df.groupby('fn').agg(
    has_ep_non_w=('docn', lambda x: (x.str.startswith('EP') & ~x.str.endswith('W')).any()),
    has_ep_w=('docn', lambda x: (x.str.startswith('EP') & x.str.endswith('W')).any())
)
# 筛选出同时满足两个条件的fn
valid_fns = fn_flags[(fn_flags['has_ep_non_w'] & fn_flags['has_ep_w'])].index
# 从原DataFrame中筛选这些fn对应的行
result_df = df[df['fn'].isin(valid_fns)]

两种方法最终都会得到符合要求的结果，包含fn=007658和fn=007690的所有行。

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