如何在Pandas中按规则保留DataFrame指定列的3个值？

给定如下Pandas DataFrame df（包含日期、实际值和多列滚动预测结果）：

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个有效数值
• 保留的起始点为2021-11-28行的20211003列，之后按从左到右、从下到上逐天递增的逻辑保留对应行的数值，最终期望结果如下：

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          NaN          NaN         NaN
4   2021-10-17  3572.3662          NaN          NaN         NaN
5   2021-10-24  3582.6036          NaN          NaN         NaN
6   2021-10-31  3547.3361          NaN          NaN         NaN
7    2021-11-7  3491.5677          NaN          NaN         NaN
8   2021-11-14  3539.1002          NaN          NaN         NaN
9   2021-11-21  3560.3734          NaN          NaN         NaN
10  2021-11-28  3564.0894  3255.328902          NaN         NaN
11   2021-12-5  3607.4320  3313.274702  3396.912886         NaN
12  2021-12-12  3666.3479  3371.220502  3450.172564  3412.23444
13  2021-12-19  3632.3638          NaN  3466.930383  3428.68349
14  2021-12-26  3618.0535          NaN          NaN  3370.73769

实现代码 #

import pandas as pd

# 将date列转为datetime类型，确保日期匹配准确
df['date'] = pd.to_datetime(df['date'])

# 定位2021-11-28对应的行索引
start_idx = df[df['date'] == '2021-11-28'].index[0]

# 提取所有预测列（value列之后的列）
prediction_cols = df.columns[2:]

# 遍历每个预测列，按规则保留指定行的数值
for col_idx, col_name in enumerate(prediction_cols):
    # 计算当前列需要保留的连续3行索引
    keep_rows = range(start_idx + col_idx, start_idx + col_idx + 3)
    # 先保存原始值，再将整列设为缺失值，最后把需要保留的行赋值回去
    original_values = df[col_name].copy()
    df[col_name] = pd.NA
    df.loc[keep_rows, col_name] = original_values.loc[keep_rows]

代码解释 #

1. 日期格式转换：把date列转为datetime类型，避免字符串匹配误差，确保精准定位起始行。
2. 起始索引定位：找到2021-11-28对应的行索引，作为保留数据的基准点。
3. 预测列提取：通过df.columns[2:]直接获取value之后的所有预测列，无需硬编码列名，提升代码通用性。
4. 逐列处理：
   • 对每个预测列，根据它在列表中的位置（col_idx）计算需要保留的行索引：基准索引加上列的偏移量，取连续3行。
   • 先保存该列的原始值，再将整列设为缺失值，最后把需要保留的行的原始值赋值回去，实现仅保留指定行的效果。

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