Stratified Split

Split the indices of a label vector into a train set and a test set so that every class keeps (approximately) the same proportion in both. Within each class the indices are shuffled with np.random.default_rng(seed) and the first int(count * test_size) of them go to the test set.

Implement solve(y, test_size=0.2, seed=0).

Returns. Return a tuple (train_idx, test_idx) of two sorted NumPy index arrays. Because of int(...) the number of test items per class is rounded down, so a class with fewer than $1\mathrm{/}\text{test\_size}1/\backslash text\{test\backslash \_size\}$ members contributes nothing to the test set.

Examples

Example 1

Input:

solve([0, 0, 0, 0, 1, 1, 1, 1], 0.5, 0)

Output:

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

Example 2

Input:

solve([0, 0, 0, 0, 0, 0, 0, 0, 1, 1], 0.25, 3)

Output:

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