Medium

Stratified Sample Allocation

~20 mincode completion

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

Examples

Equal classes: proportional allocation divides evenly

Input
stratified_sample_counts([100, 100], 10)
Output
[5, 5]

80/20 population split: allocation matches proportions exactly

Input
stratified_sample_counts([80, 20], 10)
Output
[8, 2]

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

Input
stratified_sample_counts([100, 200, 300], 10)
Output
[2, 3, 5]

Hints

Hint 1

You need the index of the extreme value, not the value itself.

Hint 2

Reach for largest remainder method rather than round.

Requirements

  • class_counts: 1D array of per-class population sizes, shape (k,)

  • total_samples: Total number of samples to allocate

  • Return Integer array of shape (k,) summing to total_samples.

Constraints

  • Allowed library: NumPy only

  • Time limit: 200 ms, Memory: 64 MB

Try similar problems(3)
Easy
Undersample for Class Balance

Dataset Splitting & Sampling · ~15 min

Easy
Walk-Forward Cross-Validation Splits

Dataset Splitting & Sampling · ~15 min

Intro
K-Fold Size Allocation

Dataset Splitting & Sampling · ~10 min

Where this shows up

~20 min

6 employers weight this skill

3 quant funds, 2 big tech firms, 1 health and bio company. Top match scores 71.

Python
import numpy as np

def stratified_sample_counts(class_counts: np.ndarray, total_samples: int) -> np.ndarray:
    """
    Compute per-class sample counts via stratified sampling (largest remainder method).

    Args:
        class_counts:  1D array of per-class population sizes, shape (k,)
        total_samples: Total number of samples to allocate

    Returns:
        Integer array of shape (k,) summing to total_samples.
    """
    # YOUR CODE HERE
    pass
Loading docs…

The AI Mentor needs an account

It reads your code and the failing tests and nudges you toward the fix without handing you the answer. Free accounts get it on every problem you're working on today.