如何拆分pandas dataframe并筛选指定年份的训练集与测试集

首先导入依赖库：
import pandas as pd

1 构造示例数据（已有数据可跳过） #

按照你描述的字段和数据构造测试用DataFrame：

data = [
    ["2017-01", "A", 100],
    ["2018-02", "A", 110],
    ["2019-01", "A", 90],
    ["2019-02", "A", 105],
    ["2020-01", "A", 117],
    ["2020-02", "A", 120]
]
df = pd.DataFrame(data, columns=["year-month", "company", "GWh"])

2 删除2020年所有数据 #

提供两种实现方式：

• 字符串前缀匹配（简单直接，适合格式统一的场景）

df = df[~df["year-month"].str.startswith("2020")]

• 转为时间格式过滤（兼容性更强，适合后续时间维度分析）

df["year-month"] = pd.to_datetime(df["year-month"])
df = df[df["year-month"].dt.year != 2020]

注意：如果选择第二种转时间格式的方案，后续拆分数据集也可以直接用时间属性判断，逻辑更严谨

3 拆分训练集、测试集 #

3.1 基于字符串前缀拆分 #

# 训练集：2019年之前的所有记录
train_df = df[df["year-month"].str.startswith(("2017", "2018"))]
# 测试集：2019年的所有记录
test_df = df[df["year-month"].str.startswith("2019")]

3.2 基于时间属性拆分（适配第二步转时间格式的方案） #

train_df = df[df["year-month"].dt.year < 2019]
test_df = df[df["year-month"].dt.year == 2019]

结果验证 #

最终得到的两个数据集内容如下：

• 训练集包含2条记录：2017-01、2018-02的电力消耗数据
• 测试集包含2条记录：2019-01、2019-02的电力消耗数据

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