Implement Stratified Sampling by Country

Problem Statement

You are given a list of dictionaries representing user interactions, a categorical column name stratify_col (e.g., "country"), and a target sample size k. Write a function to extract an unbiased, stratified sample of size k without replacement, preserving the proportion of records in each stratum as closely as possible.

Because the proportional allocation k * (count_group / total_count) typically results in non-integer values, use the Largest Remainder Method (Hare-Niemeyer method) to allocate sample counts:

1. Compute the exact float quota for each group: quota = k * (group_size / total_size).
2. Assign each group its integer floor: base_quota = floor(quota).
3. Compute each group's remainder: remainder = quota - base_quota.
4. If the sum of base quotas is less than k, distribute the remaining slots one by one to groups in descending order of their remainders. If two groups have identical remainders, break ties by sorting the stratum keys in ascending order.
5. Randomly sample each group's allocated quota without replacement using a reproducible random generator seeded with seed.

Function Signature

def stratified_sample(
    data: list[dict],
    stratify_col: str,
    k: int,
    seed: int | None = None
) -> list[dict]:
    pass

Input & Output Format

• Input:
  • data: A list of dictionaries, where each dictionary contains at least the key specified by stratify_col.
  • stratify_col: A string specifying the column to stratify by.
  • k: An integer representing the target sample size (0 <= k <= len(data)).
  • seed: An optional integer for reproducibility.
• Output:
  • A list of k dictionaries sampled proportionally from data.

Constraints

• 0 <= len(data) <= 100,000
• 0 <= k <= len(data)
• Every dictionary in data contains the key stratify_col with a non-null, hashable value.

Example 1

data = [
    {"id": 1, "country": "US"},
    {"id": 2, "country": "US"},
    {"id": 3, "country": "US"},
    {"id": 4, "country": "US"},
    {"id": 5, "country": "CA"},
    {"id": 6, "country": "CA"},
    {"id": 7, "country": "UK"},
    {"id": 8, "country": "UK"},
    {"id": 9, "country": "FR"},
    {"id": 10, "country": "FR"}
]
k = 5

Calculation:

• Total size = 10.
• "US": 5 * (4/10) = 2.0 -> base 2, remainder 0.0
• "CA": 5 * (2/10) = 1.0 -> base 1, remainder 0.0
• "UK": 5 * (2/10) = 1.0 -> base 1, remainder 0.0
• "FR": 5 * (2/10) = 1.0 -> base 1, remainder 0.0
• Total base = 5, matches k = 5. The output must contain exactly 2 records from "US", 1 from "CA", 1 from "UK", and 1 from "FR".

Example 2

data = [
    {"id": 1, "country": "US"},
    {"id": 2, "country": "US"},
    {"id": 3, "country": "CA"}
]
k = 2

Calculation:

• Total size = 3.
• "US": 2 * (2/3) = 1.333... -> base 1, remainder 0.333...
• "CA": 2 * (1/3) = 0.666... -> base 0, remainder 0.666...
• Sum of base quotas = 1. Remaining slots to allocate = 2 - 1 = 1.
• "CA" has a higher remainder (0.666... > 0.333...), so "CA" receives the extra slot.
• Final quotas: "US": 1, "CA": 1.