def swap_examples(examples_1,
                  examples_2,
                  get_compounds_fn,
                  get_atoms_fn,
                  max_iterations=1000,
                  max_divergence=None,
                  min_atom_count=1,
                  print_frequencies=True):
  """Swaps examples between examples_1 and examples_2 to maximize divergence.

  This approach first balances atoms to ensure that every atom that appears
  in examples_2 appears in examples_1 at least `min_atom_count` times.
  Then, the algorithm identifies a swap between each collection of examples
  that does not violate the atom constraint, but increases compound divergence.
  The procedure breaks when no swap that increases compound divergence can
  be found, or max_iterations or max_divergence is reached.

  To generate different splits, a different initial random split into examples_1
  and examples_2 can be used before calling this function.

  Args:
    examples_1: A list of examples of type E.
    examples_2: A list of examples of type E.
    get_compounds_fn: A function from E to a collections.Counter of strings
      representing compounds.
    get_atoms_fn: A function from E to a set of strings representing atoms.
    max_iterations: A maximum number of iterations (i.e. swap) to run for.
    max_divergence: If not None, will break if compound divergence exceeds this
      value.
    min_atom_count: Minimum amount of times an atom in examples_2 should appear
      in examples_1.
    print_frequencies: Whether to print compound frequencies at the end of
      swapping.

  Returns:
    (examples_1, examples_2) where each list is the same length and type as the
    corresponding input, but examples have been swapped per the method described
    above.
  """
  examples_1, examples_2 = balance_atoms(examples_1, examples_2, get_atoms_fn,
                                         max_iterations, min_atom_count)
  examples_1, examples_2 = maximize_divergence(examples_1, examples_2,
                                               get_compounds_fn, get_atoms_fn,
                                               max_iterations, max_divergence,
                                               min_atom_count)
  if print_frequencies:
    print_compound_frequencies(examples_1, examples_2, get_compounds_fn)
  return examples_1, examples_2
Swaps examples between examples_1 and examples_2 to maximize divergence.

  This approach first balances atoms to ensure that every atom that appears
  in examples_2 appears in examples_1 at least `min_atom_count` times.
  Then, the algorithm identifies a swap between each collection of examples
  that does not violate the atom constraint, but increases compound divergence.
  The procedure breaks when no swap that increases compound divergence can
  be found, or max_iterations or max_divergence is reached.

  To generate different splits, a different initial random split into examples_1
  and examples_2 can be used before calling this function.

  Args:
    examples_1: A list of examples of type E.
    examples_2: A list of examples of type E.
    get_compounds_fn: A function from E to a collections.Counter of strings
      representing compounds.
    get_atoms_fn: A function from E to a set of strings representing atoms.
    max_iterations: A maximum number of iterations (i.e. swap) to run for.
    max_divergence: If not None, will break if compound divergence exceeds this
      value.
    min_atom_count: Minimum amount of times an atom in examples_2 should appear
      in examples_1.
    print_frequencies: Whether to print compound frequencies at the end of
      swapping.

  Returns:
    (examples_1, examples_2) where each list is the same length and type as the
    corresponding input, but examples have been swapped per the method described
    above.
  
language/compgen/nqg/tasks/mcd_utils.py