PY-98

Split Whole Groups under a Rule

  • Medium–Hard
  • Grouped Splits
  • Python

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:

  1. 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.
  2. 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.