import collections


def _get_atoms_below_count(examples, get_atoms_fn, max_count=1):
  """Return set of atoms that appear <= max_count times across all examples."""
  # Map of atom to count of examples containing atom.
  atoms_to_count = collections.defaultdict(int)
  for example in examples:
    atoms = get_atoms_fn(example)
    for atom in atoms:
      atoms_to_count[atom] += 1

  single_atoms = set(
      [atom for atom, count in atoms_to_count.items() if count <= max_count])
  return single_atoms