Your question is Balanced Subset Sampling. Start with the requirements on the right.
Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.
Cresta model-training workflows may need a fixed number of examples across intent classes while respecting the available examples in each class. Given a requested total n and the capacity of every class, return a deterministic allocation that is as balanced as possible.
Use the following definition of balanced: maximize the smallest allocation across all classes, then distribute any remaining examples so that no eligible class receives more than one extra example over another. Classes are processed in their original order when a tie remains.
Implement balance_samples(n, capacities), where n is a nonnegative integer and capacities[i] is the number of available examples for class i. Return an array allocation of the same length such that 0 <= allocation[i] <= capacities[i] and sum(allocation) == n.
It is guaranteed that n does not exceed the total capacity. Do not sample actual examples, only compute counts.
def balance_samples(n, capacities):