def _get_swap_idx(examples, atoms, get_atoms_fn, contains=True):
  """Returns an example based on a constraint over atoms.

  If `contains` is True, returns an example in `examples` that contains any
  atom in `atoms`.

  If `contains` is False, returns an example in `examples` that does not contain
  any atom in `atoms`.

  Args:
    examples: List of examples.
    atoms: Set of atoms.
    get_atoms_fn: Function from an example to set of atoms.
    contains: Bool (see function docstring for usage).

  Returns:
    (example_idx, example) for example meeting criteria in docstring.
  """
  for example_idx, example in enumerate(examples):
    example_contains_atom = _contains_atom(example, atoms, get_atoms_fn)
    if example_contains_atom == contains:
      return example_idx, example
  if contains:
    raise ValueError("Could not find example that contains any atoms in: %s" %
                     atoms)
  else:
    raise ValueError(
        "Could not find example that doesn't contain any atoms in: %s" % atoms)
Returns an example based on a constraint over atoms.

  If `contains` is True, returns an example in `examples` that contains any
  atom in `atoms`.

  If `contains` is False, returns an example in `examples` that does not contain
  any atom in `atoms`.

  Args:
    examples: List of examples.
    atoms: Set of atoms.
    get_atoms_fn: Function from an example to set of atoms.
    contains: Bool (see function docstring for usage).

  Returns:
    (example_idx, example) for example meeting criteria in docstring.
  
language/compgen/nqg/tasks/mcd_utils.py