Как убедиться, что мой набор данных равномерно распределен по классам, т. Е. Он стратифицирован, размер и распределение классов должны быть сбалансированы?

Question

Я сделал простой код перекрестной проверки K-fold, теперь я хочу внести некоторые изменения, чтобы он был сбалансирован по размеру и распределению классов ?.

P.S .: Мне нужно использовать код Python с нуля, sklearn не допускается.

from random import seed
from random import randrange




def cross_validation_split(dataset, folds=3):
    dataset_split = []
    dataset_copy = list(dataset)
    fold_size = int(len(dataset) / folds)
    for i in range(folds):
        fold = []
        while len(fold) < fold_size:
            index = randrange(len(dataset_copy))
            fold.append(dataset_copy.pop(index))
        dataset_split.append(fold)
return dataset_split


seed()
dataset = [1,1,1,2,2,2,3,3,4,4,4]
folds = cross_validation_split(dataset, 2)
print(folds)

В результате я получаю: [[4, 4, 3, 3, 4], [1, 1, 2, 2, 1]].

Я хочу, чтобы это было, например, [[1,3,2,4,4],[1,2,2,4,3]]