如何在Pandas中按行指定列范围填充NA值为0

我有如下Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'End' : ['2022-03','2022-05','2022-06'],
    '2022-01' : [1,2,np.nan],
    '2022-02' : [np.nan,3,4],
    '2022-03' : [np.nan,1,3],
    '2022-04' : [np.nan,np.nan,2],
    '2022-05' : [np.nan,np.nan,np.nan],
    '2022-06' : [np.nan,np.nan,np.nan]
})

需要填充每行的NaN值，规则为：将每行中End列指定的列及之前的列内的NaN替换为0，End列之后的列保持NaN不变。

期望输出如下：

pd.DataFrame({
    'End' : ['2022-03','2022-05','2022-06'],
    '2022-01' : [1,2,0],
    '2022-02' : [0,3,4],
    '2022-03' : [0,1,3],
    '2022-04' : [np.nan,0,2],
    '2022-05' : [np.nan,0,0],
    '2022-06' : [np.nan,np.nan,0]
})

方法一：逐行处理（直观易懂） #

先提取所有日期列并建立列名与位置的映射，再逐行根据End列的值生成填充掩码，替换对应区域的NaN：

# 提取除End外的日期列
date_cols = [col for col in df.columns if col != 'End']
# 建立列名到索引位置的映射
col_position = {col: idx for idx, col in enumerate(date_cols)}

def process_row(row):
    end_col_name = row['End']
    end_pos = col_position[end_col_name]
    # 生成掩码：标记End列及之前的所有列
    fill_mask = [i <= end_pos for i in range(len(date_cols))]
    # 对掩码覆盖区域的NaN填充0，其余保持原样
    row[date_cols] = np.where(fill_mask, row[date_cols].fillna(0), row[date_cols])
    return row

# 应用到每行得到结果
result_df = df.apply(process_row, axis=1)

方法二：向量化操作（高效适用于大数据集） #

通过广播生成二维掩码矩阵，避免逐行循环，提升处理效率：

date_cols = df.columns.drop('End')
# 获取每行End对应的列索引位置
end_indices = df['End'].apply(date_cols.get_loc)
# 生成二维掩码：每行中列索引 <= 对应End位置的区域标记为True
mask = np.arange(len(date_cols)) <= end_indices[:, np.newaxis]
# 复制原数据并按掩码填充NaN
result_df = df.copy()
result_df[date_cols] = np.where(mask, result_df[date_cols].fillna(0), result_df[date_cols])

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