如何高效优化Pandas中基于type和locale的多条件赋值代码？

原代码中重复编写了Holiday和Event的条件判断逻辑，以下是几种更简洁高效的实现方式，结果与原函数完全一致：

方法1：映射字典+向量化操作（推荐，效率最高） #

通过先将locale映射为对应分值，再用np.where根据type筛选赋值，避免重复的条件组合：

import pandas as pd
import numpy as np

def preprocess(data_input):
    data = data_input.copy()
    # 定义locale到分值的映射关系
    locale_score_map = {'National': 3, 'Regional': 2, 'Local': 1}
    # 将locale转换为对应的分值
    data['_locale_score'] = data['locale'].map(locale_score_map)
    
    # 生成Holiday列：type为Holiday时取分值，否则为0
    data['Holiday'] = np.where(data['type'] == 'Holiday', data['_locale_score'], 0)
    # 生成Events列：type为Event时取分值，否则为0
    data['Events'] = np.where(data['type'] == 'Event', data['_locale_score'], 0)
    
    # 删除临时中间列
    data.drop('_locale_score', axis=1, inplace=True)
    return data

方法2：apply一行式赋值（写法直观） #

借助apply函数遍历每行，根据type和locale直接赋值，代码更紧凑：

import pandas as pd
import numpy as np

def preprocess(data_input):
    data = data_input.copy()
    locale_score_map = {'National': 3, 'Regional': 2, 'Local': 1}
    
    data['Holiday'] = data.apply(
        lambda row: locale_score_map[row['locale']] if row['type'] == 'Holiday' else 0,
        axis=1
    )
    data['Events'] = data.apply(
        lambda row: locale_score_map[row['locale']] if row['type'] == 'Event' else 0,
        axis=1
    )
    return data

方法3：One-Hot编码+乘法扩展（适合多类型场景） #

如果后续需要新增更多type的分值列，这种方法扩展性更强：

import pandas as pd
import numpy as np

def preprocess(data_input):
    data = data_input.copy()
    locale_score_map = {'National': 3, 'Regional': 2, 'Local': 1}
    data['_locale_score'] = data['locale'].map(locale_score_map)
    
    # 仅对需要的type生成One-Hot编码
    type_dummies = pd.get_dummies(data['type'])[['Holiday', 'Event']]
    # 编码结果与分值相乘得到最终列
    data[['Holiday', 'Events']] = type_dummies.mul(data['_locale_score'], axis=0)
    
    data.drop('_locale_score', axis=1, inplace=True)
    return data

验证结果一致性 #

可以通过以下代码确认优化后函数与原函数的输出完全一致：

# 原函数（用于对比）
def preprocess_original(data_input):
    data = data_input.copy()
    
    conditions = [
        (data['type'] ==  'Holiday') & (data.locale == 'National'),
        (data['type'] ==  'Holiday') & (data.locale == 'Regional'),
        (data['type'] ==  'Holiday') & (data.locale == 'Local')
        ]
    values = [3,2,1]
    data.loc[:,'Holiday'] = np.select(conditions, values, default=0)
    
    conditions = [
        (data['type'] ==  'Event') & (data.locale == 'National'),
        (data['type'] ==  'Event') & (data.locale == 'Regional'),
        (data['type'] ==  'Event') & (data.locale == 'Local')
        ]
    values = [3,2,1]
    data.loc[:,'Events'] = np.select(conditions, values, default=0)
    return(data)

# 测试数据
df = pd.DataFrame({'date': {0: '2012-03-02',  1: '2012-04-01',  2: '2012-04-12',  3: '2012-04-14',  4: '2012-04-21',  5: '2012-05-12',  6: '2012-06-23',  7: '2012-06-25',  8: '2012-06-26',  9: '2012-06-27'},
                   'type': {0: 'Holiday',  1: 'Other',  2: 'Event',  3: 'Holiday',  4: 'Event',  5: 'Holiday',  6: 'Other',  7: 'Holiday',  8: 'Holiday',  9: 'Event'},
                   'locale': {0: 'Local',  1: 'Regional',  2: 'National',  3: 'Local',  4: 'National',  5: 'National',  6: 'National',  7: 'Regional',  8: 'Regional',  9: 'Local'}})

# 对比结果
print(preprocess_original(df).equals(preprocess(df)))  # 输出True

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