We all know that with the use of sklearn package from python, we can create X_train, X_test, y_train and y_test via this code:

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33)

I want to make sure that each training and testing phase of a multi-class data set, have 66.33/33.33 percent of each class values so the prediction and accuracy would get better. All i want is 66.33 percent of class A in training set and 33.33 percent of Class A in test set. And, so on for other classes, like B, C, D and etc. in a given multi-class data set.

Is the code provided enough to achieve this or should i write extra code?



1 Answer 1


You need to call train_test_split(X, y, test_size=0.33, stratify=y) assuming that y contains your class labels. Alternatively, you can call StratifiedShuffleSplit directly.

  • $\begingroup$ Thanks. Sorry, that it doesn't let me vote your answer. ./ $\endgroup$
    – Jack Kay
    Commented Dec 11, 2018 at 4:52

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.