如何创建函数基于Pandas DataFrame多列取值生成符合规则的新列

首先我们把分类规则整合成一个统一的嵌套映射字典，方便后续匹配调用：

import pandas as pd

# 统一规则字典：一级key为犬只type，二级结构为{颜色: 对应breed列表/单个值}
colour_rule = {
    'toy': {
        'black' : ['poodle', 'shih_tzu'], 
        'mixed' : 'king_charles_spaniel'
    },
    'terrier': {
        'black_brown' : ['yorkshire_terrier','bull_terrier'],
        'white' : 'west_highland',
        'white_orange' : 'fox_terrier'
    },
    'hound': {
        'brindle' : 'greyhound',
        'brown' : 'afghan'
    }
}

然后编写可复用的匹配函数：

def add_colour_column(df, type_col='type', breed_col='breed', rule=colour_rule, default_val=None):
    """
    根据type和breed为DataFrame新增colour列
    参数：
        df: 待处理的DataFrame
        type_col: df中存储犬只type的列名，默认'type'
        breed_col: df中存储犬只breed的列名，默认'breed'
        rule: 颜色匹配规则字典，结构为{type: {colour: [breed_list]/breed_str}}
        default_val: 未匹配到规则时的默认返回值，默认None
    返回：
        新增colour列后的DataFrame
    """
    def match_colour(row):
        current_type = row[type_col]
        current_breed = row[breed_col]
        # 先匹配type，不存在直接返回默认值
        if current_type not in rule:
            return default_val
        type_rule = rule[current_type]
        # 遍历当前type下的颜色规则匹配breed
        for colour, breed_info in type_rule.items():
            # 统一转成列表判断，兼容单个breed字符串和多个breed列表两种格式
            if current_breed in ([breed_info] if isinstance(breed_info, str) else breed_info):
                return colour
        return default_val
    
    df['colour'] = df.apply(match_colour, axis=1)
    return df

使用示例 #

用测试DataFrame调用函数：

# 构造测试DataFrame
df = pd.DataFrame({
    'type': ['terrier', 'toy','toy','toy', 'hound' , 'terrier', 'terrier', 'terrier','terrier', 'hound'],
    'breed' : ['yorkshire_terrier', 'king_charles_spaniel', 'poodle', 'shih_tzu', 'greyhound', 'west_highland', 'bull_terrier' , 'fox_terrier', 'west_highland', 'afghan']
})

# 调用函数
df_res = add_colour_column(df)
print(df_res)

输出结果和预期完全一致：

type                 breed        colour
0  terrier     yorkshire_terrier   black_brown
1      toy  king_charles_spaniel         mixed
2      toy                poodle         black
3      toy              shih_tzu         black
4    hound             greyhound       brindle
5  terrier         west_highland         white
6  terrier          bull_terrier   black_brown
7  terrier           fox_terrier  white_orange
8  terrier         west_highland         white
9    hound                afghan         brown

方案特性 #

• 复用性高：支持自定义列名、规则字典、未匹配默认值，可直接套用在其他结构相似的DataFrame上
• 匹配准确：必须同时匹配type和breed才会返回对应颜色，不会出现跨type的错误匹配
• 兼容性强：规则字典支持单个breed字符串、多个breed列表两种格式，不需要调整原有规则结构

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