1 / 1

Stratified Train/Test Split

Data Processing

Hard

Split features and labels while approximating the same class proportions in train and test. For each class with $n_c$ samples, compute:

n_c^{\mathrm{test}} = \operatorname{round}(n_c\,t)

Here, $t$ is test_size. When a class has more than one sample, cap its test count at $n_c-1$ so at least one remains in training. Shuffle each class with np.random.default_rng(seed), combine the selected indices, sort each final index set, and return X_train, X_test, y_train, and y_test as NumPy arrays in a dictionary.

Loading visualization...

Examples

Input: X = [0, 1, 2, 3, 4, 5], y = [0, 0, 0, 1, 1, 1], test_size = 0.33, seed = 42

Output: {"X_train": [0, 1, 4, 5], "X_test": [2, 3], "y_train": [0, 0, 1, 1], "y_test": [0, 1]}

Explanation: One seeded sample from each class enters the test split.

Input: X = [[1, 0], [2, 0], [3, 0], [4, 0], [5, 0], [6, 0], [7, 0], [8, 0], [9, 0], [10, 0]], y = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1], test_size = 0.3, seed = 42

Output: {"X_train": [[1, 0], [2, 0], [5, 0], [6, 0], [7, 0], [8, 0], [9, 0]], "X_test": [[3, 0], [4, 0], [10, 0]], "y_train": [0, 0, 0, 0, 0, 1, 1], "y_test": [0, 0, 1]}

Hint 1

Use np.flatnonzero(y == label) and rng.permutation for each class.

Hint 2

Accumulate class indices, then apply np.sort before indexing X and y.

Requirements

Shuffle indices independently within each class using seed
Round each class test count and preserve one training sample when possible
Sort the combined train and test indices before indexing
Return exactly X_train, X_test, y_train, and y_test in a dictionary of NumPy arrays

Constraints

X and y contain the same number of samples
y is a nonempty one-dimensional label list
test_size is greater than 0 and less than 1
Use NumPy only

Prev Next