Я предлагаю использовать train_test_split
с опцией stratify
, а test_size
= 0,5 следующим образом:
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
np.random.seed(123)
df = pd.DataFrame({"Age": np.random.randint(0,50,10000),
"work_exp" : np.random.randint(0,3,10000),
"man_exp" : np.random.randint(0,3,10000),
"value": np.random.randint(0,2,10000)})
df_train, df_test = train_test_split(df, test_size=0.5,
stratify=df[["Age", "work_exp", "man_exp"]], random_state=0)
Вот результат для основной статистики как df_train
, так и df_test
:
> df_train.describe()
Age work_exp man_exp value
count 5000.000000 5000.000000 5000.00000 5000.000000
mean 24.516600 1.014200 1.01520 0.493400
std 14.453107 0.820812 0.81431 0.500006
min 0.000000 0.000000 0.00000 0.000000
25% 12.000000 0.000000 0.00000 0.000000
50% 24.000000 1.000000 1.00000 0.000000
75% 37.000000 2.000000 2.00000 1.000000
max 49.000000 2.000000 2.00000 1.000000
> df_test.describe()
Age work_exp man_exp value
count 5000.00000 5000.000000 5000.000000 5000.00000
mean 24.47900 1.011200 1.009400 0.51000
std 14.45663 0.819762 0.815503 0.49995
min 0.00000 0.000000 0.000000 0.00000
25% 12.00000 0.000000 0.000000 0.00000
50% 24.00000 1.000000 1.000000 1.00000
75% 37.00000 2.000000 2.000000 1.00000
max 49.00000 2.000000 2.000000 1.00000
Надеюсь, это поможет