Class-Balanced Sampler Indices

Problem

Build a class-balanced index sequence.

Group indices by label, then cycle through classes in ascending order, appending the next index of each class and wrapping within the class when it runs out, until num_samples indices are produced.

Hint

Use a per-class counter with modulo so exhausted classes wrap around.

Test cases

case 1

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

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

case 2

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

expected: [0,2,3]

case 3

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

expected: []

case 4

in: [[],5]

expected: []

Starter codePython

def balanced_indices(labels, num_samples):
    # Your code here
    pass

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