def balance_atoms(examples_1, examples_2, get_atoms_fn, max_iterations,
                  min_atom_count):
  """Attempts to ensure every atom is represented in the first set."""
  for iteration_num in range(max_iterations):
    # Find atoms that appear >= min_atom_count in `examples_1`.
    atoms_1_above = _get_atoms_above_count(examples_1, get_atoms_fn,
                                           min_atom_count)

    # Find atoms in `examples_2`.
    atoms_2 = get_all_atoms(examples_2, get_atoms_fn=get_atoms_fn)

    # Find atoms in examples_2 not in examples_1 at least `min_atom_count`.
    atoms_2_m_1 = atoms_2 - atoms_1_above

    # If there are no atoms in `atoms_2` not in  in examples_1, then we have
    # reached our goal state.
    if not atoms_2_m_1:
      print("Atoms are balanced after %s iterations." % iteration_num)
      return examples_1, examples_2

    # Find atoms that appear <= min_atom_count in `examples_1`.
    atoms_1_below = _get_atoms_below_count(examples_1, get_atoms_fn,
                                           min_atom_count)

    # Find candidates to swap.
    # First, find an example in examples_1 that does not contain any atoms
    # that appear only once in examples_1. Otherwise, moving the examples to
    # examples_2 can take us farther from our goal state.
    example_1_idx, example_1 = _get_swap_idx(
        examples_1, atoms_1_below, get_atoms_fn=get_atoms_fn, contains=False)
    # Second, find an example in examples_2 that contains one of the atoms
    # that is currently missing from examples_1.
    example_2_idx, example_2 = _get_swap_idx(
        examples_2, atoms_2_m_1, get_atoms_fn=get_atoms_fn, contains=True)

    # Swap the examples.
    del examples_1[example_1_idx]
    examples_1.append(example_2)
    del examples_2[example_2_idx]
    examples_2.append(example_1)

  raise ValueError("Could not find split that balances atoms [%s] [%s]" %
                   (atoms_1_below, atoms_2_m_1))
Attempts to ensure every atom is represented in the first set.
language/compgen/nqg/tasks/mcd_utils.py