如何在Python Pandas中按70%/30%唯一ID拆分DataFrame并保留全量行

输入数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "id": [999, 999, 999, 185, 185, 185, 44, 44, 44],
    "target": [1, 1, 1, 0, 0, 0, 1, 1, 1],
    "event_date": ['2023-01-01', '2023-01-02', '2023-01-03', 
                   '2023-01-01', '2023-01-02', '2023-01-03', 
                   '2023-01-01', '2023-01-02', '2023-01-03'],
    "event1": [1, 6, 11, 16, np.nan, 22, 74, 109, 52],
    "event2": [2, 7, np.nan, 17, 22, np.nan, np.nan, 10, 5],
    "event3": [3, 8, 13, 18, 23, np.nan, 2, np.nan, 99],
    "event4": [4, 9, np.nan, np.nan, np.nan, 11, 8, np.nan, np.nan],
    "event5": [5, np.nan, 15, 20, 25, 1, 1, 3, np.nan]
})

# 用0填充缺失值
df = df.fillna(0)

需求 #

• 训练集包含70%的唯一ID，测试集包含剩余30%的唯一ID
• 每个ID对应的所有行必须完整保留，不能拆分单个ID的记录

解决方案 #

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

# 随机打乱ID顺序（保证拆分随机性）
np.random.shuffle(unique_ids)

# 计算拆分阈值，取70%的ID作为训练集
split_point = int(len(unique_ids) * 0.7)
train_ids = unique_ids[:split_point]
test_ids = unique_ids[split_point:]

# 按ID筛选生成训练集和测试集
train_df = df[df['id'].isin(train_ids)].reset_index(drop=True)
test_df = df[df['id'].isin(test_ids)].reset_index(drop=True)

结果示例 #

• 训练集（示例，包含2个ID，真实场景自动按70%比例分配）：

# 训练集数据（已填充0）
train_df = pd.DataFrame({
    "id": [999, 999, 999, 185, 185, 185],
    "target": [1, 1, 1, 0, 0, 0],
    "event_date": ['2023-01-01', '2023-01-02', '2023-01-03', 
                   '2023-01-01', '2023-01-02', '2023-01-03'],
    "event1": [1, 6, 11, 16, 0, 22],
    "event2": [2, 7, 0, 17, 22, 0],
    "event3": [3, 8, 13, 18, 23, 0],
    "event4": [4, 9, 0, 0, 0, 11],
    "event5": [5, 0, 15, 20, 25, 1]
})

• 测试集（示例，包含1个ID）：

# 测试集数据（已填充0）
test_df = pd.DataFrame({
    "id": [44, 44, 44],
    "target": [1, 1, 1],
    "event_date": ['2023-01-01', '2023-01-02', '2023-01-03'],
    "event1": [74, 109, 52],
    "event2": [0, 10, 5],
    "event3": [2, 0, 99],
    "event4": [8, 0, 0],
    "event5": [1, 3, 0]
})

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