stratified-split — Booking.com case

Booking.com-inspired experimentation pipeline needs training and validation sets whose class proportions remain stable when the dataset is split. You need to construct a deterministic stratified split so rare outcome classes are represented in both partitions.

Split the sample indices into training and validation sets so that each class is split in (nearly) the same proportion. For every class, the indices of that class are shuffled with np.random.default_rng(seed) (classes are visited in sorted order, one generator for all of them), and round(count * val_fraction) of them go to validation. Python's round rounds halves to the nearest even integer.

Implement solve(y,val_fraction,seed=0).

Returns. Return a tuple (train_idx, val_idx) of sorted NumPy index arrays. A class so small that its rounded validation count is 0 contributes only to training.

Examples

Example 1

Input:

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

Output:

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

Example 2

Input:

solve([0, 0, 0, 0, 1, 1], 0.5, seed=3)

Output:

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