import random


def randomize_strata(n_items, group_ids, seed=None):
  """Perform stratified randomization.
  Maps a number of items into group ids by dividing the items into strata of
  length equal to the number of groups, then assigning the group ids randomly
  within each strata. If the number of items do not divide the number of groups
  equally, the remaining items are assigned to groups randomly (with equal
  probability of assignment in each group).
  Example: randomize_strata(12, [1, 2, 3]) yields a list of 12 group ids, each
  1, 2, or 3. There are (3!)^4 = 6^4 = 1296 possible outcomes as there are 12 /
  3 = 4 strata, and within each strata there are 3! = 6 different ways to assign
  the items into groups. Therefore, the first 3 items of the list (that is,
  positions 0, 1, and 2) can be mapped only to the 3!=6 different mappings:
  [1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2] or [3, 2, 1].
  Args:
    n_items: (int) Number of items to assign to groups.
    group_ids: A list of group ids that are typically integers or strings, but
      can be of any type.
    seed: (int) Random seed, applied to a local instance of class random.Random;
      if not specified, the global random instance is used instead.
  Returns:
    A list of length n_items, consisting of the group ids whose positions in the
    list correspond to the items.
  """

  if seed is None:
    random_sampler = random.sample
  else:
    random_sampler = random.Random(seed).sample

  n_groups = len(group_ids)
  n_strata = n_items // n_groups
  groups = []
  for _ in range(n_strata):
    groups.extend(random_sampler(group_ids, n_groups))
  remaining = n_items - len(groups)
  if remaining > 0:
    groups.extend(random_sampler(group_ids, remaining))
  return groups