Rebalance Experiment Buckets with Minimal Reassignment

Problem Statement

In experiment infrastructure, entity allocations across variants (buckets) frequently need adjustment as experiment sample sizes change. To preserve user experience and avoid invalidating experiment tracking, you must rebalance assignments to exact target counts while retaining as many existing bucket assignments as possible.

Write a function rebalance_buckets(current_assignments: dict[str, str], target_counts: dict[str, int]) -> dict[str, str] that takes:

• current_assignments: a mapping of entity IDs to their current bucket names.
• target_counts: a mapping of bucket names to their desired entity counts.

Your function must return a new dictionary containing the updated assignments that satisfy the target counts and adhere to the following rules:

1. Minimal Reassignments: An entity must not be reassigned if it can stay in its current bucket without causing the bucket to exceed its target count.
2. Canonical Retention: For any bucket whose current count exceeds its target count, deterministically select which entities stay by retaining entities with the lexicographically smallest entity IDs. The remaining entities are released into an unassigned pool.
3. Canonical Reassignment: Sort the displaced entities lexicographically by entity ID. Reassign them to deficit buckets (buckets where current count < target count) considered in lexicographical order of their bucket names.
4. Non-mutation: Do not mutate the input dictionaries.

Function Signature

def rebalance_buckets(current_assignments: dict[str, str], target_counts: dict[str, int]) -> dict[str, str]:
    pass

Constraints

• 1 <= len(current_assignments) <= 10^5
• 1 <= len(target_counts) <= 10^3
• sum(target_counts.values()) == len(current_assignments)
• All entity IDs and bucket names are non-empty strings containing alphanumeric characters and underscores.
• All bucket names in target_counts have non-negative target counts (target_counts[b] >= 0).

Examples

Example 1

current = {
    "user_1": "control",
    "user_2": "control",
    "user_3": "control",
    "user_4": "treatment"
}
targets = {
    "control": 2,
    "treatment": 2
}

result = rebalance_buckets(current, targets)
# Output: {"user_1": "control", "user_2": "control", "user_3": "treatment", "user_4": "treatment"}

Explanation: "control" has 3 users and a target of 2. "user_1" and "user_2" are retained because they are lexicographically smallest. "user_3" is displaced and moved to "treatment" to fulfill its deficit of 1.

Example 2

current = {
    "uA": "bucket_1",
    "uB": "bucket_2"
}
targets = {
    "bucket_2": 1,
    "bucket_3": 1
}

result = rebalance_buckets(current, targets)
# Output: {"uA": "bucket_3", "uB": "bucket_2"}

Explanation: bucket_1 is not in the targets (implicit target of 0). "uA" is displaced and reassigned to the deficit bucket "bucket_3".