Python实现多类别自动生成时间序列拆分的训练测试索引

要解决这个动态适配任意类别数量的问题，我们可以利用字典来存储不同类别的测试集索引，替代手动创建每个类别的变量。这样不管你的数据里有多少种cat值，代码都能自动处理，完全实现参数化运行。

优化后的代码实现 #

import pandas as pd

# 示例数据
df = pd.DataFrame({
    'date': [1, 1, 2, 2, 2, 3, 3, 3, 4, 4],
    'prod': ['a', 'b', 'a', 'b', 'c', 'a', 'b', 'c', 'b', 'd'],
    'cat': ['c1', 'c2', 'c1', 'c2', 'c2', 'c1', 'c2', 'c2', 'c2', 'c1']
})

time_split = [2, 3, 4, 5]

# 初始化存储容器
train_index = []
test_index = []
# 自动识别所有类别，用字典存储每个类别的测试集索引
category_test_indices = {cat: [] for cat in df['cat'].unique()}

# 遍历每个时间拆分窗口
for i in range(len(time_split) - 1):
    # 生成当前拆分的训练集索引
    train_idx = df.loc[df['date'] < time_split[i]].index
    # 生成当前拆分的整体测试集索引
    test_idx = df.loc[(df['date'] >= time_split[i]) & (df['date'] < time_split[i+1])].index
    
    # 为每个类别筛选对应的测试集索引
    for cat in category_test_indices.keys():
        cat_test_idx = test_idx[df.loc[test_idx, 'cat'] == cat]
        category_test_indices[cat].append(cat_test_idx)
    
    # 存入整体训练、测试索引
    train_index.append(train_idx)
    test_index.append(test_idx)

代码核心思路 #

1. 自动识别类别：通过df['cat'].unique()获取数据中所有唯一的类别，动态生成字典的键，彻底告别硬编码c1、c2等固定类别的操作。
2. 字典统一存储：category_test_indices字典的每个键对应一个类别，值是该类别在每个时间拆分下的测试集索引列表，结构和train_index、test_index完全对齐，方便后续调用。
3. 循环批量处理：在时间拆分的循环中嵌套遍历所有类别，自动筛选每个类别在当前测试窗口中的索引，避免重复编写大量类似idx3、idx4的冗余代码。

结果验证 #

比如查看c1类别的测试索引：

print(category_test_indices['c1'])
# 输出：
# [Int64Index([2], dtype='int64'), Int64Index([5], dtype='int64'), Int64Index([9], dtype='int64')]

这个结果和你手动代码中cat1_index的输出完全一致。如果后续数据新增c3、c4等类别，代码会自动为这些类别生成对应的索引列表，完全不需要修改代码结构。

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