Neural Mastery
← Practice
Stratified Train-Test Split Partitioneasy

Stratified Train-Test Split Partition

easy NumPy & Data Wrangling 10–15 min

NumPy allowed · plain Python earns 10 bonus points

Linear algebra step 1/4

A vector transforms

Matrix multiplication is the first computation behind every model.
Open lesson →
01vector
02matrix
03transform
04result
[2, −1] × [[1, 2], [−.5, 1]] = [0, −2]

Computed matrix × vector transformation

Partition dataset indices into train and test sets while preserving class label proportions.

Task

Implement stratified_split_indices(labels, test_size). Group indices by label (processed in alphabetical order of labels). For each label with indices C, calculate n_test = round(len(C) * test_size) and n_train = len(C) - n_test. Put the first n_train indices into train and the remaining n_test into test. Return (sorted(train_indices), sorted(test_indices)).

Function Signature

stratified_split_indices(labels: list[str], test_size: float) -> tuple[list[int], list[int]]

Examples

Example 1: Balanced binary labels
Input: labels = ['cat', 'dog', 'cat', 'dog', 'cat', 'dog'] test_size = 0.33
Output: [[0, 1, 2, 3], [4, 5]]
Example 2: Imbalanced classes
Input: labels = ['a', 'a', 'a', 'a', 'b', 'b'] test_size = 0.5
Output: [[0, 1, 4], [2, 3, 5]]

Constraints

  • labels is a list of strings
  • 0.0 < test_size < 1.0
Python3Saved ✓
def stratified_split_indices(labels: list[str], test_size: float) -> tuple[list[int], list[int]]:
"""Split sample indices into train and test preserving class balance.
Returns (train_indices, test_indices) both sorted ascending.
"""
pass

Case 1: Balanced binary labels
Input: labels = ['cat', 'dog', 'cat', 'dog', 'cat', 'dog'] test_size = 0.33
Expected: [[0, 1, 2, 3], [4, 5]]
Case 2: Imbalanced classes
Input: labels = ['a', 'a', 'a', 'a', 'b', 'b'] test_size = 0.5
Expected: [[0, 1, 4], [2, 3, 5]]