Python中如何实现按分组而非随机划分的k折交叉验证

你需要的拆分逻辑核心是保证同一个个体的全生命周期数据不会被拆分到训练集和测试集，以下是两种可直接落地的实现方案：

方案1：使用scikit-learn内置GroupKFold（推荐） #

GroupKFold是scikit-learn专门为分组拆分场景设计的交叉验证接口，原生支持按指定分组完整划分，不会出现同组数据跨折的问题，实现代码如下：

from sklearn.model_selection import GroupKFold
import pandas as pd
import numpy as np

# 你的已有代码
grouped = data.groupby('id')
lifetimes_list = grouped.size()

# 准备分组标识、特征、标签
groups = data['id'].values
X = data.drop('你的目标列名', axis=1) # 替换为你的特征列
y = data['你的目标列名'] # 替换为你的标签列

# 初始化K折，n_splits替换为你需要的折数
kf = GroupKFold(n_splits=5)

# 遍历每折的训练、测试索引
for train_idx, test_idx in kf.split(X, y, groups=groups):
    X_train, X_test = X.iloc[train_idx], X.iloc[test_idx]
    y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]
    # 此处添加你的模型训练、评估逻辑

方案2：手动实现自定义分组拆分 #

如果你需要自定义拆分逻辑（比如按个体寿命分层、控制折间样本量均衡），可以基于唯一分组ID手动拆分，代码如下：

# 固定随机种子保证结果可复现
np.random.seed(42)

# 获取所有独立个体的ID列表
unique_ids = list(grouped.groups.keys())
# 打乱ID顺序
np.random.shuffle(unique_ids)

# 设置折数
k = 5
# 将ID列表拆分为k份
id_folds = np.array_split(unique_ids, k)

for fold_idx in range(k):
    # 取当前折的测试集ID，剩余ID为训练集
    test_ids = id_folds[fold_idx]
    train_ids = np.concatenate(id_folds[:fold_idx] + id_folds[fold_idx+1:])
    
    # 按ID筛选训练、测试集
    train_data = data[data['id'].isin(train_ids)]
    test_data = data[data['id'].isin(test_ids)]
    # 此处添加你的模型训练、评估逻辑

如果需要保证每个折的总观测数尽量均衡，可以基于你提前计算的lifetimes_list做加权分层拆分，避免部分折样本量偏差过大。

注意事项 #

• 两种方案都完全避免同一个体的数据同时出现在训练集和测试集导致的数据泄露问题，符合你的需求
• 若你不需要对拆分逻辑做特殊定制，优先使用GroupKFold方案，代码更简洁且经过工业级验证
• 若数据量较小，可以适当降低折数，避免单个折内样本量不足影响评估结果的可靠性

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