Stratified Train/Test Split
Stratified Train/Test Split
Split features X and labels y into train/test while preserving class proportions (stratification).
Stratified splitting maintains the original class distribution in both training and test sets. For each class, calculate how many samples should go to the test set based on the test_size ratio, then randomly select those samples while ensuring at least one sample per class remains in the training set when possible.
Function Arguments
X: array-like, shape (N, D) or (N,)- Featuresy: array-like, shape (N,)- Integer class labelstest_size: float in (0,1)- Fraction for test setrng: optional np.random.Generator- For reproducible shuffling
Examples
Input: X = [0,1,2,3,4,5], y = [0,0,0,1,1,1], test_size = 0.33
Output: X_train (4 items), X_test (2 items), y_train, y_test
Class 0 (3 samples): round(3 × 0.33) = 1 in test. Class 1 (3 samples): 1 in test. Both train and test preserve the 50/50 class ratio.
Input: X = [[1,0],...,[10,0]], y = [0,0,0,0,0,0,0,1,1,1], test_size = 0.3
Output: X_train (7 items), X_test (3 items)
Class 0 (7 samples): round(7 × 0.3) = 2 in test. Class 1 (3 samples): round(3 × 0.3) = 1 in test. Original ratio 70/30 is preserved in both splits.
Hint 1
Use np.unique() to get class counts, then calculate test samples per class.
Hint 2
For each class, find indices with np.where(), shuffle them, then split.
Hint 3
Use rng.shuffle() if rng provided, else np.random.shuffle().
Requirements
- Return tuple: (X_train, X_test, y_train, y_test)
- Preserve class ratios as closely as possible
- Shuffle within each class using rng if provided
- Handle minority classes and rounding robustly
- No sklearn; NumPy only
Constraints
- N ≥ 2, at least one sample per class must remain in train if possible
- NumPy only; time limit: 300ms
Log in to take notes on this problem
Accepts: array
Accepts: array
Accepts: number
Accepts: any
Accepts: number
Stratified Train/Test Split
Stratified Train/Test Split
Split features X and labels y into train/test while preserving class proportions (stratification).
Stratified splitting maintains the original class distribution in both training and test sets. For each class, calculate how many samples should go to the test set based on the test_size ratio, then randomly select those samples while ensuring at least one sample per class remains in the training set when possible.
Function Arguments
X: array-like, shape (N, D) or (N,)- Featuresy: array-like, shape (N,)- Integer class labelstest_size: float in (0,1)- Fraction for test setrng: optional np.random.Generator- For reproducible shuffling
Examples
Input: X = [0,1,2,3,4,5], y = [0,0,0,1,1,1], test_size = 0.33
Output: X_train (4 items), X_test (2 items), y_train, y_test
Class 0 (3 samples): round(3 × 0.33) = 1 in test. Class 1 (3 samples): 1 in test. Both train and test preserve the 50/50 class ratio.
Input: X = [[1,0],...,[10,0]], y = [0,0,0,0,0,0,0,1,1,1], test_size = 0.3
Output: X_train (7 items), X_test (3 items)
Class 0 (7 samples): round(7 × 0.3) = 2 in test. Class 1 (3 samples): round(3 × 0.3) = 1 in test. Original ratio 70/30 is preserved in both splits.
Hint 1
Use np.unique() to get class counts, then calculate test samples per class.
Hint 2
For each class, find indices with np.where(), shuffle them, then split.
Hint 3
Use rng.shuffle() if rng provided, else np.random.shuffle().
Requirements
- Return tuple: (X_train, X_test, y_train, y_test)
- Preserve class ratios as closely as possible
- Shuffle within each class using rng if provided
- Handle minority classes and rounding robustly
- No sklearn; NumPy only
Constraints
- N ≥ 2, at least one sample per class must remain in train if possible
- NumPy only; time limit: 300ms
Log in to take notes on this problem
Accepts: array
Accepts: array
Accepts: number
Accepts: any
Accepts: number