如何在Pandas中迭代多列，将指定行后的值替换为NaN？

需求说明 #

给定如下原始DataFrame：

date      value     20211003     20211010     20211017
0    2021-9-19  3613.9663          NaN          NaN          NaN
1    2021-9-26  3613.0673          NaN          NaN          NaN
2    2021-10-3  3568.1668          NaN          NaN          NaN
3   2021-10-10  3592.1666  3510.221000          NaN          NaN
4   2021-10-17  3572.3662  3465.737012  3534.220800          NaN
5   2021-10-24  3582.6036  3479.107035  3539.856801  3514.420400
6   2021-10-31  3547.3361  3421.161235  3481.911001  3456.474600
7    2021-11-7  3491.5677  3370.140147  3439.284539  3416.621024
8   2021-11-14  3539.1002  3319.289523  3391.930037  3370.079953
9   2021-11-21  3560.3734  3261.343723  3333.984237  3312.134153
10  2021-11-28  3564.0894  3255.328902  3338.967086  3305.054247
11   2021-12-5  3607.4320  3313.274702  3396.912886  3363.000047
12  2021-12-12  3666.3479  3371.220502  3450.172564  3412.234440
13  2021-12-19  3632.3638          NaN  3466.930383  3428.683490
14  2021-12-26  3618.0535          NaN          NaN  3370.737690

需要对value列之后的所有预测列（20211003、20211010、20211017）执行以下操作：

• 仅保留从首次出现有效值的行开始的连续3个有效值
• 将该列后续所有行的值替换为NaN
• 最终得到如下结果：

date      value     20211003     20211010     20211017
0   2021-9-19  3613.9663          NaN          NaN          NaN
1   2021-9-26  3613.0673          NaN          NaN          NaN
2   2021-10-3  3568.1668          NaN          NaN          NaN
3  2021-10-10  3592.1666  3510.221000          NaN          NaN
4  2021-10-17  3572.3662  3465.737012  3534.220800          NaN
5  2021-10-24  3582.6036  3479.107035  3539.856801  3514.420400
6  2021-10-31  3547.3361          NaN  3481.911001  3456.474600
7   2021-11-7  3491.5677          NaN          NaN  3416.621024

Pandas实现方案 #

可以通过遍历目标列，定位有效行范围并替换超出范围的值来实现：

import pandas as pd

# 假设原始DataFrame已加载为df
# 筛选出需要处理的预测列（跳过date和value）
prediction_cols = df.columns[2:]

for col in prediction_cols:
    # 找到当前列第一个非NaN值的索引
    first_valid_pos = df[col].first_valid_index()
    if first_valid_pos is not None:
        # 计算需要保留的最后一行索引（连续3行：起始、起始+1、起始+2）
        last_valid_pos = first_valid_pos + 2
        # 将起始+3及以后的行对应列设为NaN
        df.loc[last_valid_pos + 1:, col] = pd.NA

# 可选：如果需要截断到所有预测列都无有效值的行之前（匹配示例结果）
# 找到所有预测列最后一个有效值的最大索引
max_last_valid = max(df[col].last_valid_index() for col in prediction_cols)
df = df.loc[:max_last_valid]

代码说明 #

1. 筛选目标列：通过df.columns[2:]获取所有预测列，避开前两列的日期和原始值。
2. 定位有效行起始点：使用first_valid_index()找到每个预测列第一个非空值的位置。
3. 确定保留范围：从起始点开始连续3行，因此最后保留的行索引是起始点+2。
4. 替换超出范围的值：用loc选中超出范围的行，将对应列的值设为pd.NA。
5. 可选截断DataFrame：如果需要和示例结果一样只保留到有有效值的最后一行，找到所有预测列最后一个有效值的最大索引，截断DataFrame到该行。

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