如何用for循环简化Polars中含super promo类型的when-then表达式函数

原函数针对每个week_offset（取值1-3），依次判断三种super promo类型（TV、YouTube、Banners），匹配时取对应促销均值列，否则取无超级促销的均值列。你尝试的写法仅处理了TV的条件，且循环生成的列名重复，未关联week_offset，因此无法正常运行。

正确的实现思路是：对每个week_offset，先初始化when表达式，再通过循环依次追加三种super promo的判断条件，最后补充otherwise分支，完整代码如下：

def add_promo_means(data: pl.DataFrame) -> pl.DataFrame:
    # 统一管理super promo类型，后续增减只需修改此列表
    superpromo_types = ['TV', 'YouTube', 'Banners']
    columns = []
    
    for week_offset in range(1, 4):
        # 初始化第一个super promo的判断逻辑
        expr = pl.when(pl.col(f"is_{superpromo_types[0]}_superpromo") == 1)\
                 .then(pl.col(f"mean_promo_quantity_with_{superpromo_types[0]}_0_{week_offset}_promo_weeks"))
        
        # 循环追加剩余super promo的判断条件
        for promo in superpromo_types[1:]:
            expr = expr.when(pl.col(f"is_{promo}_superpromo") == 1)\
                       .then(pl.col(f"mean_promo_quantity_with_{promo}_0_{week_offset}_promo_weeks"))
        
        # 补充默认分支并设置列别名
        expr = expr.otherwise(pl.col(f"mean_promo_quantity_without_superpromo_0_{week_offset}_promo_weeks"))\
                   .alias(f"mean_promo_sales_0_{week_offset}_promo_weeks")
        columns.append(expr)
    
    return data.with_columns(columns)

核心说明： #

• 将super promo类型抽象为列表，后续修改类型无需改动判断逻辑，提升可维护性
• 外层循环处理每个week_offset，确保生成的列与原函数输出完全对应
• 内层循环批量追加多条件判断，消除原函数中重复的when-then代码块

内容的提问来源于stack exchange，提问作者Татьяна Вахрушева