Dividing data set into (almost) equal batches based on characteristics I have a large dataset containing 6 characteristics (all numerical).
I need to split this dataset into multiple batches to be processed in parallel, and ideally, the batches should be as equal in size as possible.
The catch is: I can only split based on the 6 characteristics, so I have to specify ranges of values for each batch.
Simplified example: [1,1,1,2,2,3,3,3,3,4] will be split in two by specifying numbers 1-2 go to batch 1 and 3-4 go to batch two and I end up with [1,1,1,2,2] and [3,3,3,3,4].
I managed to make a simple division algorithm based on one characteristic, unfortunately some of the characteristics might not be suitable to split on, and I have no way to reliably predict which characteristic will create a good split.
I remember from Uni that a classification algorithm might do the trick here, but I am not sure how to implement the requirement for similar size classes.
here is the simple code I wrote for one characteristic, it may help to understand the situation:
def divide_into_groups(lst: list, num_groups=2):
    length = len(lst)
    sorted_list = sorted(lst)
    optimal_group_size = length // num_groups

    initial_group_markers = [sorted_list[i * optimal_group_size] for i in range(1, num_groups)]
    marker_edges = [(sorted_list.index(marker), length - 1 - sorted_list[::-1].index(marker))
                    for marker in initial_group_markers]

    final_group_markers = []
    for i, marker in enumerate(initial_group_markers):
        if i - marker_edges[i][0] &lt; marker_edges[i][1] - i:
            final_group_markers.append(marker - 1)
        else:
            final_group_markers.append(marker)

    distances = [abs(initial_group_markers[i] - final_group_markers[i]) for i in range(len(initial_group_markers))]
    efficiency = (length - sum(distances)) / length
    print('distances from optimals: %s' % distances)
    print('efficiency of division is %s %%' % (efficiency * 100))

    return final_group_markers

P.S. this is my first question in this forum. If I am missing some information or have made any mistake, please comment and I will fix as soon as possible.
If I got this right you want to split the dataset into train and test sets in a way that preserves the same proportions of examples in each class as observed in the original dataset?
This is called a stratified train-test split. See the stratify argument here sklearn split