scikit_learn/2_stratify.py at main · fayzi-dev/scikit_learn · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import pandas as pd

df = pd.DataFrame({'feture':list(range(12)), 'target':['No']*10 + ['Yes']*2})
# print(df)
X = df[['feture']]
y = df['target']


# print(X)

from sklearn.model_selection import train_test_split

#Not Startify
X_train, X_test, y_train, y_test = train_test_split(X,y, test_size=0.5, random_state=0)
print(y_train)
#output
# 1    No
# 7    No
# 9    No
# 3    No
# 0    No
# 5    No
print(y_test)
#output
# 6      No
# 11    Yes
# 4      No
# 10    Yes
# 2      No
# 8      No


# Startify
X_train, X_test, y_train, y_test = train_test_split(X,y, test_size=0.5, random_state=0, stratify=y)
print(y_train)
# Output
# 9      No
# 8      No
# 4      No
# 1      No
# 11    Yes
# 2      No
print(y_test)
# Output
# 10    Yes
# 3      No
# 0      No
# 5      No
# 7      No
# 6      No