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:
labels = [0, 0, 0, 0, 1, 1, 1, 1, 1, 1] train_ratio = 0.5
([0, 1, 4, 5, 6], [2, 3, 7, 8, 9])
- 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
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.