如何筛选Pandas中所有行status均为red的id？

给定如下Pandas DataFrame，需要找出所有对应行的status字段全为"red"的id：

import pandas as pd

data = {"id":[1,1,1,2,2,2],
        "status":["red","red","red","green","red","red"]}
df = pd.DataFrame(data)

原代码因布尔序列索引与原DataFrame索引不匹配，触发IndexingError错误。以下是几种可行的解决方法：

方法1：分组后直接校验全量匹配 #

通过groupby对id分组，检查每组内所有status是否都等于"red"，再提取符合条件的id：

# 分组校验后筛选出符合条件的id并转为列表
all_rows_are_red = df.groupby('id')['status'].apply(lambda x: (x == 'red').all())
all_rows_are_red = all_rows_are_red[all_rows_are_red].index.tolist()

print(all_rows_are_red)  # 输出: [1]

方法2：反向排除法 #

先找出所有出现过非"red"状态的id，再从所有唯一id中排除这些值：

# 获取所有存在非red状态的id
invalid_ids = df[df['status'] != 'red']['id'].unique()
# 从所有唯一id中排除无效id
all_rows_are_red = [id_val for id_val in df['id'].unique() if id_val not in invalid_ids]

print(all_rows_are_red)  # 输出: [1]

方法3：用transform标记每行状态 #

通过transform给每行添加一个标记，指示该id对应的所有行是否全为"red"，再提取唯一的有效id：

# 给每行添加标记：该id是否所有行都是red
df['is_all_red'] = df.groupby('id')['status'].transform(lambda x: (x == 'red').all())
# 筛选标记为True的行，提取唯一id
all_rows_are_red = df[df['is_all_red']]['id'].unique().tolist()

print(all_rows_are_red)  # 输出: [1]

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