如何在Python中根据字典值范围为DataFrame添加对应键列

问题背景 #

现有如下示例DataFrame：

import pandas as pd
df = pd.DataFrame({'Price':[10,8,7,6,10,12,11,11,7,9], 'Group':['apple','apple','apple','apple','apple','berry','berry','berry','berry','berry']})

对应的表格：

另有按Group分组的规则字典：

apple = {'A':[9, 10], 'B':[6, 8], 'C':[3,5]}
berry = {'A':[11, 12], 'B':[6, 9]}
# 合并为总规则字典
group_rules = {'apple': apple, 'berry': berry}

需求：为DataFrame添加Cat列，当Price处于对应Group字典中某个键的闭区间数值范围时，填入该键，预期输出如下：

实现方案 #

可以通过自定义函数结合pandas方法实现，以下提供两种可行方案：

方案1：逐行匹配规则（小数据量适用） #

代码简洁直观，直接遍历每行数据匹配对应分组的规则：

def get_category(row, rules):
    # 获取当前行的分组规则与价格
    current_rules = rules[row['Group']]
    price = row['Price']
    # 遍历规则，返回匹配的分类
    for cat, (low, high) in current_rules.items():
        if low <= price <= high:
            return cat
    # 无匹配时返回None（可根据需求修改）
    return None

# 调用函数添加Cat列
df['Cat'] = df.apply(get_category, axis=1, rules=group_rules)

# 查看结果
print(df)

方案2：向量化匹配（大数据量适用） #

利用pandas的向量化操作提升效率，避免逐行遍历：

# 将规则转换为结构化DataFrame
rule_records = []
for group, cat_rules in group_rules.items():
    for cat, (low, high) in cat_rules.items():
        rule_records.append({'Group': group, 'Cat': cat, 'low': low, 'high': high})
rule_df = pd.DataFrame(rule_records)

# 合并原数据与规则表，筛选符合价格区间的记录
merged_data = df.merge(rule_df, on='Group', how='left')
matched_data = merged_data[merged_data['Price'].between(merged_data['low'], merged_data['high'])]

# 按原数据索引分组，取第一个匹配的分类
category_result = matched_data.groupby(matched_data.index)['Cat'].first().reset_index()

# 合并回原DataFrame
df = df.merge(category_result, left_index=True, right_on='index').drop('index', axis=1)

print(df)

两种方案均能得到预期结果，可根据数据量大小选择合适的实现方式。

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