Stratified Train-Test Split

mediumFREEModel EvaluationMachine Learning 1

Implement stratified train-test split that preserves the class distribution in both sets.

Given data indices and their labels, split into train and test sets such that each set has approximately the same proportion of each class.

For each unique label, assign floor(count * train_ratio) samples to training and the rest to test (using the original order of indices within each class).

Return a tuple (train_indices, test_indices) where both lists are sorted.

Example

Input

labels = [0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
train_ratio = 0.5

Output

([0, 1, 4, 5, 6], [2, 3, 7, 8, 9])

Reasoning

labels gives one class per data point; position $i$ in the list is data point $i$ . Here points $0$ – $3$ are class 0 and points $4$ – $9$ are class 1 — a $40\%$ / $60\%$ mix.
train_ratio $= 0.5$ is applied to each class separately. Class 0: $\lfloor 4 \times 0.5 \rfloor = 2$ points go to training. Class 1: $\lfloor 6 \times 0.5 \rfloor = 3$ points go to training.
Within a class, the earliest points are the training ones: points $0, 1$ from class 0 and $4, 5, 6$ from class 1. The rest ( $2, 3$ and $7, 8, 9$ ) becomes the test set.
Both sets are reported in ascending order, giving ([0, 1, 4, 5, 6], [2, 3, 7, 8, 9]).
Notice the training set is $2$ of class 0 and $3$ of class 1: still $40\%$ / $60\%$ , the same mix as the full data. Because of the floor, an odd class size (e.g. $5$ at ratio $0.5$ ) sends the extra point to test.

Constraints

labels: list of class labels (integers)
train_ratio: float between 0 and 1
Return tuple of two sorted lists of indices
Preserve class proportions as closely as possible

Learn this in

Introduction to MLML Study Plan ClassificationML Study Plan Model EvaluationML Study Plan

More from Machine Learning 1

RBF Kernel Matrixhard Polynomial Kernelmedium Hinge Losseasy Linear Decision Boundaryeasy Gradient Boosting Residual Stephard

Concept primer

Background Knowledge

The stratified train-test split is a technique used in machine learning to split a dataset into training and testing sets while preserving the class distribution. This is particularly important when dealing with imbalanced datasets, where some classes have a significantly larger number of instances than others. By maintaining the same proportion of each class in both the training and testing sets, we can ensure that our model is not biased towards the majority class and that the evaluation metrics are more accurate.

In model evaluation, it's essential to have a representative test set that reflects the real-world scenario. A stratified split helps to achieve this by ensuring that the test set has the same class distribution as the overall dataset. This is in contrast to a random split, where the class distribution in the test set may be different from the overall dataset, leading to biased evaluation metrics. The train ratio determines the proportion of samples assigned to the training set, with the remaining samples assigned to the test set.

The concept of stratification is not unique to machine learning and is used in various fields, such as statistics and survey research. In the context of machine learning, stratification is used to split the data into subsets based on the target variable, ensuring that each subset has the same proportion of each class. This technique is particularly useful when dealing with classification problems, where the goal is to predict a categorical label.

Algorithm/Approach

The general approach to solving this problem involves the following algorithm pattern:

Group the data indices by their corresponding labels
For each unique label, calculate the number of samples to assign to the training set based on the train ratio
Split the samples for each label into training and test sets, using the calculated number of training samples
Combine the training and test sets from all labels, sorting the resulting lists of indices

Step-by-Step Strategy

To implement the solution, follow these steps:

Create a dictionary to store the data indices for each unique label
Iterate over the data indices and labels, grouping the indices by their corresponding labels
For each unique label, calculate the number of samples to assign to the training set using the train ratio
Split the samples for each label into training and test sets, using the calculated number of training samples
Combine the training and test sets from all labels, sorting the resulting lists of indices

Common Pitfalls

When implementing the solution, watch out for the following:

Ensure that the train ratio is a value between 0 and 1, where 0 means all samples are assigned to the test set and 1 means all samples are assigned to the training set
Use the floor function to calculate the number of training samples for each label, to ensure that the number of samples is an integer
Sort the resulting lists of indices to ensure that the output is in the correct order

Time & Space Complexity

The expected time complexity for this solution is O(n log n) due to the sorting of the resulting lists of indices, where n is the total number of data indices. The space complexity is O(n), as we need to store the data indices for each unique label. Note that the space complexity can be reduced to O(k), where k is the number of unique labels, if we only store the indices for each label and not the entire dataset.

solution.py

Python 3.13

Test results