如何在Python Pandas中按两列唯一组合拆分DataFrame为70%/30%数据集

输入数据 #

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': ['2023-01-01', '2023-01-02', '2023-02-03', '2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03'],
    'survey': ['2023-02-02', '2023-02-02', '2023-02-02', '2023-03-10', '2023-03-10', '2023-03-10', '2023-04-22', '2023-04-22', '2023-04-22'],
    '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]
})

df = df.fillna(0)

需求说明 #

• 按id和survey的唯一组合拆分数据集，70%组合对应全部行作为训练集
• 剩余30%组合对应全部行作为测试集
• 每个组合的所有行必须完整保留，不能拆分部分行

解决方法 #

步骤1：提取唯一的ID-Survey组合 #

先获取所有不重复的id+survey组合：

unique_pairs = df[['id', 'survey']].drop_duplicates().reset_index(drop=True)

步骤2：拆分组合为训练和测试集 #

用train_test_split按比例拆分这些唯一组合，确保拆分的是组合而非单一行：

from sklearn.model_selection import train_test_split

train_pairs, test_pairs = train_test_split(unique_pairs, test_size=0.3, random_state=42)

random_state用于固定拆分结果，可根据需要调整

步骤3：根据拆分后的组合筛选原数据集 #

通过合并或布尔索引获取训练集和测试集：

方法一：合并筛选

# 训练集：保留属于train_pairs的所有行
train_df = df.merge(train_pairs, on=['id', 'survey'], how='inner')

# 测试集：保留属于test_pairs的所有行
test_df = df.merge(test_pairs, on=['id', 'survey'], how='inner')

方法二：布尔索引筛选（效率更高）

train_mask = df[['id', 'survey']].apply(tuple, axis=1).isin(train_pairs.apply(tuple, axis=1))
train_df = df[train_mask]

test_mask = df[['id', 'survey']].apply(tuple, axis=1).isin(test_pairs.apply(tuple, axis=1))
test_df = df[test_mask]

预期结果 #

• 训练集：包含70%的id-survey唯一组合对应的全部行
• 测试集：包含剩余30%的id-survey唯一组合对应的全部行

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