Stratified Train Test Split

Split row indices into train and test sets so that each class keeps the same share in both.

With a rare class, a random split can put almost none of it in the test set, and then the test score says nothing about that class. A stratified split avoids that by splitting each class on its own.

For each class c, in increasing order of label:

1. List the indices i with y[i] == c, in their original order.
2. Let n_test = floor(test_frac * n_c + 0.5), which rounds half up.
3. The first n_test of those indices go to the test set, the rest to train.

Return (train_idx, test_idx), each sorted in increasing order.

Function signature

def stratified_split(y, test_frac):

Example

Input

y = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 1])
stratified_split(y, test_frac=0.25)

Output

(array([1, 2, 3, 6, 7, 8, 9]), array([0, 4, 5]))

Reasoning

Class 0 has 4 rows, and 0.25 * 4 = 1, so its first row, index 0, goes to test. Class 1 has 6 rows, and 0.25 * 6 = 1.5 rounds up to 2, so indices 4 and 5 go to test. Everything else is train.

Constraints and notes

• y holds integer labels and 0 <= test_frac <= 1.
• This version is deterministic so that it can be tested. In practice, shuffle the rows with a fixed seed first.
• Return integer arrays.