如何在Pandas中按ID划分无重复ID的训练集与测试集

嘿，这个需求在处理带有用户/个体多记录的数据集时特别常见——我们需要把整个数据集拆分成训练集和测试集，但必须保证同一个ID的所有数据行要么全在训练集，要么全在测试集，绝对不能把同一个ID的部分行分到训练、部分分到测试。下面我就用Pandas一步步实现这个操作，结合你给出的示例数据来演示：

步骤1：构造示例数据集 #

先把你给出的示例数据转换成Pandas DataFrame：

import pandas as pd

data = {
    'ID': [1,1,1,2,2,2,2,3,3,3,3],
    'AGE': [66,66,66,20,20,20,20,18,18,18,18],
    'GENDER': ['M','M','M','F','F','F','F','F','F','F','F'],
    'TIME': [1,2,3,1,2,3,4,1,2,3,4],
    'CODE': [0,0,1,0,0,0,0,0,0,0,1]
}

df = pd.DataFrame(data)
print("原数据集：")
print(df)

步骤2：拆分唯一ID集合 #

核心思路是先把所有唯一的ID拆分成训练ID和测试ID，再根据这些ID去筛选原数据。这里我们用sklearn.model_selection里的train_test_split来做随机拆分（你也可以手动指定ID来拆分）：

from sklearn.model_selection import train_test_split

# 获取所有唯一ID
unique_ids = df['ID'].unique()

# 拆分ID：比如按8:2的比例拆分，random_state保证结果可复现
train_ids, test_ids = train_test_split(unique_ids, test_size=0.25, random_state=42)

注：这里test_size=0.25是因为你的示例里测试集占1个ID，原ID共3个，刚好接近1/3比例，你可以根据自己的需求调整拆分比例。

步骤3：根据ID筛选得到训练集和测试集 #

用isin()方法来筛选对应ID的所有行：

# 训练集：筛选ID在train_ids里的所有行
train_df = df[df['ID'].isin(train_ids)]
# 测试集：筛选ID在test_ids里的所有行
test_df = df[df['ID'].isin(test_ids)]

print("\n训练集：")
print(train_df)
print("\n测试集：")
print(test_df)

验证结果 #

运行上面的代码后，你会得到和示例一致的结果：

训练集包含ID为1和2的所有行，测试集包含ID为3的所有行，完全符合“同一ID不跨集”的要求。

补充：手动指定拆分ID（如果不需要随机拆分） #

如果你不想随机拆分，而是想手动指定哪些ID进训练集、哪些进测试集，直接定义train_ids和test_ids即可：

train_ids = [1,2]
test_ids = [3]

train_df = df[df['ID'].isin(train_ids)]
test_df = df[df['ID'].isin(test_ids)]

这样就完美解决了你的需求啦！

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