PY-98
Split Whole Groups under a Rule
Task
Write split_whole_groups(records, target_sizes, generator). records is a
non-empty list of dictionaries. Every record must contain exactly id and
group; both values are non-empty strings. Record IDs are unique. A group is
the set of records with the same group value, and at least three distinct
groups must be present. A group is indivisible: all of its records must go to
one partition.
target_sizes is a dictionary with exactly these keys and positive integer
values:
train, validation, test
The three targets must add to the number of records. They are desired record counts, not permission to split a group.
generator is one passed np.random.Generator. First sort the distinct group
names lexicographically. Call generator.permutation(number_of_groups) once
to obtain the group order, then use that order for all assignments. Do not
create a generator, set global random state, or make another random draw.
Use the fixed partition priority train, validation, test. Process groups
in the generated order. For each group, use this complete assignment rule:
- If one or more partitions can receive the whole group without exceeding
its target, choose the partition with the largest remaining capacity
(
target - assigned). If capacities tie, choose the first partition in the fixed priority order. - If no partition can receive the group without exceeding its target, choose
the partition with the smallest resulting overflow
(
assigned + group_size - target). If overflow ties, again choose the first partition in the fixed priority order.
After assigning a group, never move it. Within each partition, return records in generated-group order and preserve the original input order within a group. Return a dictionary with exactly these keys:
group_order, partition_groups, partitions, actual_sizes, audit
group_order is a tuple of every group name in generated order.
partition_groups maps each partition to the tuple of groups assigned to it.
partitions maps each partition to a tuple of record IDs in the order just
described. actual_sizes maps each partition to its actual record count.
audit must be a new dictionary with exactly these keys:
target_sizes, partition_order, groups_intact, all_records_once,
target_reached
target_sizes repeats the requested targets, partition_order is the fixed
tuple ("train", "validation", "test"), and the last three values are
booleans. target_reached is true only when every actual count equals its
target. Whole-group assignment may make that false; do not split or discard a
group to force a match.
Example
The group order and assignment follow the supplied generator and the stated rule. Repeating the call with a fresh generator made with the same seed gives the same order and partitions, while every input record still appears once.
Your implementation
You may import NumPy as np. Do not print or ask for input.