📘

Stratified Train-Test Split

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:

First, we separate the indices by their labels: label 0 has indices [0, 1, 2, 3] and label 1 has indices [4, 5, 6, 7, 8, 9]
Then, we calculate the number of samples to assign to the training set for each label: for label 0, $floor(4 * 0.5) = floor(2) = 2$ samples and for label 1, $floor(6 * 0.5) = floor(3) = 3$ samples
Next, we assign the calculated number of samples to the training set for each label, preserving the original order: for label 0, indices [0, 1] and for label 1, indices [4, 5, 6]
The final output is a tuple of sorted train and test indices: ([0, 1, 4, 5, 6], [2, 3, 7, 8, 9])

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

Editor

Python 3.13.1

Test Results

0/0

Run code to see test results.

📘