Medium

Stratified Sample Allocation

Medium

~20 min

code completion

Stratified Sample Allocation

Stratified sampling preserves the class distribution when selecting a subset. If class $c$ makes up fraction $f_{c}$ of the full dataset, it should make up fraction $f_{c}$ of the sample.

Because counts must be integers summing to total_samples exactly, use the largest remainder method:

1. Compute exact allocations: $raw_{c} = f_{c} \times total_samples$

2. Take the floor: $floor_{c} = ⌊ raw_{c} ⌋$

3. Compute $r = total_samples - \sum floor_{c}$ (remaining samples)

4. Assign 1 extra to the $r$ classes with the largest fractional parts

Example: class_counts = [80, 20], total_samples = 10

Fractions: [0.8, 0.2] → raw: [8.0, 2.0] → floor: [8, 2] → remainder = 0

Result: [8, 2]

Your task:

Implement stratified_sample_counts(class_counts, total_samples) that returns an integer NumPy array of per-class sample counts summing to total_samples.

Example Tests

Equal classes: proportional allocation divides evenly

Input: {"class_counts":[100,100],"total_samples":10}

Expected: [5,5]

80/20 population split: allocation matches proportions exactly

Input: {"class_counts":[80,20],"total_samples":10}

Expected: [8,2]

1/6 remainder assigned to smallest class by largest fractional part

Input: {"class_counts":[100,200,300],"total_samples":10}

Expected: [2,3,5]

You can read the full problem statement above. Create a free account to run code in the browser, submit solutions, and track your progress.