Python中如何对分组DataFrame随机排序并按玩家ID拆分训练测试集

1. 按player_id分组随机排序 #

要让每个玩家的两行数据保持关联，同时整体随机排列玩家顺序，可按以下方式实现：

import pandas as pd
import numpy as np

df = pd.DataFrame({
               "player_id":[1,1,2,2,3,3,4,4,5,5,6,6],
               "year"     :[1,2,1,2,1,2,1,2,1,2,1,2],
               "overall"  :[20,16,7,3,8,80,20,12,9,3,2,1]})

# 提取唯一玩家ID并打乱顺序
shuffled_players = np.random.permutation(df['player_id'].unique())
# 按打乱后的玩家顺序重新排列整个DataFrame
df_shuffled = df.set_index('player_id').loc[shuffled_players].reset_index()

最终df_shuffled会保持每个玩家的两行数据绑定，且玩家整体顺序随机。

2. 按80-20比例拆分训练集和测试集（无共享player_id） #

核心是先对玩家ID做分层拆分，再筛选对应数据：

# 获取所有唯一玩家ID
all_players = df['player_id'].unique()
# 计算训练集玩家数量（取整）
train_player_count = int(len(all_players) * 0.8)
# 随机选择训练集玩家
train_players = np.random.choice(all_players, size=train_player_count, replace=False)
# 拆分数据集
train_df = df_shuffled[df_shuffled['player_id'].isin(train_players)]
test_df = df_shuffled[~df_shuffled['player_id'].isin(train_players)]

这样拆分后的train_df和test_df不会有重复的玩家ID，比例接近80-20。

如果想简化成一步操作：

all_players = df['player_id'].unique()
np.random.shuffle(all_players)
split_point = int(len(all_players)*0.8)
train_players, test_players = all_players[:split_point], all_players[split_point:]

train_df = df[df['player_id'].isin(train_players)].set_index('player_id').loc[train_players].reset_index()
test_df = df[df['player_id'].isin(test_players)].set_index('player_id').loc[test_players].reset_index()

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