Stratified Split by Label

Problem

Split indices into train and validation while preserving class proportions.

Seed the random module once, then for each class in ascending order shuffle its indices and move int(round(count * val_ratio)) of them to validation. Keep the shuffled order within each split and return [train, val].

Hint

Seed once at the start and shuffle each class's index list in turn.

Test cases

case 1

in: [[0,0,0,1,1,2,2],0.5,3]

expected: [[0,4,6],[1,2,3,5]]

case 2

in: [[0,1,0,1],0.5,7]

expected: [[2,1],[0,3]]

case 3

in: [[5,5,5],0,1]

expected: [[1,2,0],[]]

Starter codePython

def stratified_split(labels, val_ratio, seed):
    # Your code here
    pass

3 test cases · run your code against them in the editor