import operator


def trim_vocab(word_frequency, vocab_size):
  """Given the max vocab size n, trim the word_frequency dictionary to only contain the
  top n occurring words
  Args:
    word_frequency (Dictionary): dictionary of word, frequency pairs.
    vocab_size (Integer): the maximum number of vocabulary allowed.
  Returns:
    retval (Dictionary): dictionary containing the top n occurring words as keys
  """
  sorted_vocab = sorted(
      word_frequency.items(), key=operator.itemgetter(1), reverse=True)
  max_count = min(len(word_frequency), vocab_size)
  retval = [k for k, _ in sorted_vocab[:max_count]]
  return retval