Pandas按条件分组保留行：实现特定规则的数据筛选

原始数据 #

import pandas as pd
df = pd.DataFrame({
    "review_num": [2,2,2,1,1,1,1,1,3],
    "review": ["The second review","The second review","The second review",
               "This is the first review","This is the first review",
               "This is the first review","This is the first review",
               "This is the first review",'No'],
    "token_num":[1,2,3,1,2,3,4,5,1],
    "token":["The","second","review","This","is","the","first","review","No"],
    "score":[0.3,-0.6,0.4,0.5,0.8,-0.7,0.6,0.4,0.3]
})

数据表格预览：

筛选规则 #

• 若review包含"t"或"r"：获取该review下，token含t或r的行中分数最高的那一行
• 若review不包含"t"或"r"：仅保留该review的唯一一行
• 保持原表中的review顺序

现有代码问题 #

现有代码仅满足规则1和3，无法捕获规则2的情况（比如review_num=3的行）：

prelist=df['token'].str.contains('|'.join(['t','r']))
token_max_score = df[prelist].groupby('review_num', sort=False)['score'].idxmax()

当前输出结果缺少review_num=3的条目：

review_num
2    2
1    6

解决方案 #

分两步处理两类规则的情况，再合并结果并保持原顺序：

# 标记review是否包含t/r（不区分大小写）
df['has_tr'] = df['review'].str.contains('[tr]', case=False)
# 标记token是否包含t/r（不区分大小写）
df['token_has_tr'] = df['token'].str.contains('[tr]', case=False)

# 处理规则1：筛选review含tr且token含tr的行，按review_num取score最大的行索引
rule1_idx = df[df['has_tr'] & df['token_has_tr']].groupby('review_num', sort=False)['score'].idxmax()

# 处理规则2：筛选review不含tr的行，取每个review_num的唯一一行索引
rule2_idx = df[~df['has_tr']].groupby('review_num', sort=False).head(1).index

# 合并索引并按原表顺序排序
final_idx = rule1_idx.append(pd.Index(rule2_idx)).sort_values()

# 获取最终的索引对应关系（与预期格式一致）
result = final_idx.rename('review_num').reset_index(drop=True)
result.index = df.loc[final_idx, 'review_num'].values

最终结果 #

执行后得到符合要求的输出：

review_num
2    2
1    6
3    8

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