Stratified Train/Test Split

Problem

Build a deterministic stratified train/test split. For each label (processed in sorted order), let idx be its positions in the original order; the first int(len(idx) * test_fraction) positions go to the test set and the rest go to train.

Return [train_indices, test_indices], both sorted in ascending order.

Hint

Per class, floor the test count and take indices from the front of that class's position list.

Test cases

case 1

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

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

case 2

in: [[0,1,0,1,0],0.4]

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

case 3

in: [[2,2,1,1,1],0.34]

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

case 4

in: [["a","b","a","b"],0.5]

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

Starter codePython

def stratified_split_indices(labels, test_fraction):
    # Returns [train_indices, test_indices]
    # Your code here
    pass

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