I would like to know how I can split in an equal number the following

0    1586
1     318

in order to have the same proportion of 0 and 1 classes in a dataset to train, if my dataset is called df and includes 10 columns, both numerical and categorical.

I would consider the following

X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.1, stratify=y)

so to do a stratification, but I do not know if it is right and I would appreciate if you could confirm it or provide an alternative to do that.

Sample of data

Fin                 Eco   Target
High percentage     12      1
Low percentage      5       0
Medium percentage   48      0
NA                  3       1
TBC                 NA      1
Low percentage      25      0
Medium percentage   12      0

How can I check if it is actually splitting in equal classes proportion my dataset? I think the best way to train a model should be having an equal proportion of 0 and 1 values. Right now I have 5 times data with Target=0.


3 Answers 3


If what you want is to have the same proportion of classes, 50% - 0 and 50% - 1. Then, there are two techniques oversampling (sampling more data from the smaller class) and undersampling (sampling less data from the bigger class). But I won't recommend you this for your problem. Your label seems fairly well balanced 1/5 it's a great proportion.

Still, this library has some implementations to do that

In this blog, you can see an overview of imbalanced datasets, but yours is not. Choosing a proper metric is more important.

If what you want to do is keep the same proportions of the label across the splits, what you are doing is right.

To validate your model properly, the class distribution and the different splits (train, validation, test) should be similar.

In the train test split documentation , you can find the argument:

stratifyarray-like, default=None If not None, data is split in a stratified fashion, using this as the class labels.

One step beyond will be using Stratified K-Folds cross-validator.

This cross-validation object is a variation of KFold that returns stratified folds. The folds are made by preserving the percentage of samples for each class.

There are more splitting techniques in Scikit Learn that you can use, have a look.

To test if the function is doing what you want just calculate the percentages in the splits:

 np.unique(y_train, return_counts=True)
 np.unique(y_val, return_counts=True)

But this will make you have the same proportions across the whole data, if your original label proportion is 1/5, then you will have 1/5 in train and 1/5 in test.

  • 2
    $\begingroup$ I kindly suggest you read the question more closely. OP asks if what they are doing "is actually splitting in equal classes proportion" their dataset; what they need to hear is that stratify=y will not do that (it will simply maintain the proportions existing in y to the produced datasets y_train and y_val), and arguably not that what they are doing is "right". $\endgroup$
    – desertnaut
    Commented Oct 11, 2020 at 22:29
  • 1
    $\begingroup$ Thanks @desertnaut, great appreciation. I edited $\endgroup$ Commented Oct 12, 2020 at 6:23

you can try stratified sampling method

from sklearn.model_selection import StratifiedShuffleSplit
split=StratifiedShuffleSplit(n_split=1, test_size=0.2, random_state=9)

This is the function I am using. You can adapt it to your needs:

# Returns a Test dataset that contains an equal amounts of each class
# y should contain only two classes 0 and 1
def TrainSplitEqualBinary(X, y, samples_n): #samples_n per class
    indicesClass1 = []
    indicesClass2 = []
    for i in range(0, len(y)):
        if y[i] == 0 and len(indicesClass1) < samples_n:
        elif y[i] == 1 and len(indicesClass2) < samples_n:
        if len(indicesClass1) == samples_n and len(indicesClass2) == samples_n:
    X_test_class1 = X[indicesClass1]
    X_test_class2 = X[indicesClass2]
    X_test = np.concatenate((X_test_class1,X_test_class2), axis=0)
    #remove x_test from X
    X_train = np.delete(X, indicesClass1 + indicesClass2, axis=0)
    Y_test_class1 = y[indicesClass1]
    Y_test_class2 = y[indicesClass2]
    y_test = np.concatenate((Y_test_class1,Y_test_class2), axis=0)
    #remove y_test from y
    y_train = np.delete(y, indicesClass1 + indicesClass2, axis=0)
    if (X_test.shape[0] != 2 * samples_n or y_test.shape[0] != 2 * samples_n):
        raise Exception("Problem with split 1!")
    if (X_train.shape[0] + X_test.shape[0] != X.shape[0] or y_train.shape[0] + y_test.shape[0] != y.shape[0]):
        raise Exception("Problem with split 2!")
    return X_train, X_test, y_train, y_test

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.