如何通过字典映射的类别键值提取DataFrame目标行？

我有一个包含CATEGORY列的DataFrame，该列值为各类农产品（如['apple', 'orange', 'banana', 'carrot', 'potato', ...]）。需要提取CATEGORY对应类别为水果而非蔬菜的行，同时有一个字典用于映射具体农产品到大类别：

produce2category = {'apple': 'fruit',
                    'banana': 'fruit',
                    'carrot': 'vegetable',
                    ...}

当前用for循环实现，但想改用更高效的方法：

>>> import pandas as pd
>>> fruit_dict = {'apple': 'fruit', 'banana': 'fruit', 'carrot': 'vegetable', 'potato': 'vegetable'}
>>> rows = [['apple', 1, 2, 3], ['potato', 4, 5, 6], ['banana', 7, 8, 9], ['carrot', 10, 11, 12]]
>>> df = pd.DataFrame(rows, columns=['CATEGORY', '1', '2', '3'])
>>> fruit_rows = []
>>> for idx, row in df.iterrows():
...     if fruit_dict[row['CATEGORY']] == 'fruit':
...         fruit_rows.append(row)
...
>>> fruit_df = pd.DataFrame(fruit_rows)

尝试用df.loc[fruit_dict[df['CATEGORY']] == 'fruit']时出现TypeError: unhashable type: 'Series'错误。

解决方案 #

方法1：使用map()函数映射后筛选 #

利用Series.map()逐个处理CATEGORY列的元素，生成对应大类别后通过布尔索引筛选：

# 直接映射并筛选
fruit_df = df[df['CATEGORY'].map(fruit_dict) == 'fruit']

# 或用loc更明确指定行列
fruit_df = df.loc[df['CATEGORY'].map(fruit_dict) == 'fruit']

方法2：提取水果列表后用isin()匹配 #

先从字典中提取所有属于fruit的农产品名称，再用isin()匹配筛选，适合重复筛选场景：

# 获取所有水果名称
fruit_names = [k for k, v in fruit_dict.items() if v == 'fruit']
# 筛选符合条件的行
fruit_df = df[df['CATEGORY'].isin(fruit_names)]

报错原因说明 #

fruit_dict[df['CATEGORY']]试图将整个Series作为键查询字典，但字典的键必须是可哈希类型（Series属于不可哈希对象），因此触发类型错误。map()方法会逐个处理Series中的每个元素，规避了这个问题。

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