Python中对不同列组循环应用where生成crosstab的更高效写法

优化后代码 #

import pandas as pd
import numpy as np 

df = pd.DataFrame({'a': ['a','b','b','a','b','b','a','b','a'], 
                   'month':['march','march','january', 'march','january','january', 'may','march','march'],
                   'week':['1','1','1', '1','2','3', '3','2','1'],
                   'year':['5','3','4', '3','1','1', '1','1','1']})

# 定义需要替换小于2取值的列
cols_need_replace = ['month', 'week']
# 所有要生成交叉表的列
target_cols = cols_need_replace + ['year']
out = {}

for col in target_cols:
    # 只生成一次交叉表，避免重复计算
    ct = pd.crosstab(df['a'], df[col])
    # 判断是否需要替换值
    if col in cols_need_replace:
        ct = ct.where(ct >= 2, 'group_a')
    out[col] = ct

优化说明 #

• 合并了原代码的两次循环、两个中间字典，只需要一次循环就能生成所有结果，减少冗余代码和不必要的内存占用
• 去掉了无意义的pd.DataFrame()包裹：pd.crosstab返回结果本身就是DataFrame类型，where方法返回的也是同结构的DataFrame，不需要重复转换
• 后续维护更方便：如果需要新增/删除需要替换值的列，只需要修改cols_need_replace列表即可，不需要调整循环逻辑
• 最终输出的out字典和你原代码的out3完全一致，符合需求。

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