Stratified K-Fold Fold Sizes

Problem

Compute the sample count of every fold in a stratified k-fold split. For each class in sorted label order, fold f receives count // k samples plus one extra when f is among the first (count % k) folds, rotating the remainder start by the class position so remainder samples spread across folds. Return the list of k fold sizes; k <= 0 returns [].

Hint

Every class is split across folds as evenly as possible, so each fold gets a similar class mix.

Test cases

case 1

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

expected: [3,3]

case 2

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

expected: [4,2]

case 3

in: [[1,1],3]

expected: [1,1,0]

case 4

in: [[],2]

expected: [0,0]

case 5

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

expected: [1,2,1]

Starter codePython

def stratified_fold_sizes(labels, k):
    # Your code here
    pass

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