Pandas如何按每行start/end范围给非区间数值列赋值NaN

通用向量化解决方案（兼容普通/跨年周数场景） #

不需要遍历行列，完全基于numpy广播实现，性能远高于迭代方案，同时支持普通连续周数和跨年周数场景：

import pandas as pd
import numpy as np

def mask_out_of_range(df):
    # 提取周数列，排除start、end字段，列名转整数方便比较
    week_cols = df.columns.difference(['start', 'end']).astype(int)
    # 把周数、start、end转成支持广播的数组格式
    weeks_arr = week_cols.to_numpy()[np.newaxis, :]
    start_arr = df['start'].to_numpy()[:, np.newaxis]
    end_arr = df['end'].to_numpy()[:, np.newaxis]
    
    # 生成保留掩码：跨年场景范围是>=start 或 <=end，非跨年是>=start 且 <=end
    mask = np.where(
        start_arr > end_arr,
        (weeks_arr >= start_arr) | (weeks_arr <= end_arr),
        (weeks_arr >= start_arr) & (weeks_arr <= end_arr)
    )
    
    # 掩码外的位置赋值为NaN
    df[week_cols.astype(str)] = df[week_cols.astype(str)].where(mask)
    return df

测试效果

对普通场景的df调用：

df = pd.DataFrame({'39' : [1, np.nan, 3],
                   '40' : [2, 4, 5],
                   '41' : [3, 1, 4],
                   '42' : [2, 5, 2],
                   '43' : [1, 1, np.nan],
                   'start' : [39, 40, 41],
                   'end' : [41, 41, 43]})
print(mask_out_of_range(df))

输出完全符合预期。
对跨年场景的df2调用：

df2 = pd.DataFrame({'51' : [1, np.nan, 3],
                   '52' : [2, 4, 5],
                   '1' : [3, 1, 4],
                   '2' : [2, 5, 2],
                   '3' : [1, 1, 3],
                   'start' : [51, 52, 52],
                   'end' : [1, 2, 1]})
print(mask_out_of_range(df2))

输出和预期结果一致。

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