Top K Frequent Words

Given a composition with different kinds of words, return a list of the top K most frequent words in the composition.

Assumptions

the composition is not null and is not guaranteed to be sorted
K >= 1 and K could be larger than the number of distinct words in the composition, in this case, just return all the distinct words

Return

a list of words ordered from most frequent one to least frequent one (the list could be of size K or smaller than K)

Examples

Composition = ["a", "a", "b", "b", "b", "b", "c", "c", "c", "d"], top 2 frequent words are [“b”, “c”]
Composition = ["a", "a", "b", "b", "b", "b", "c", "c", "c", "d"], top 4 frequent words are [“b”, “c”, "a", "d"]
Composition = ["a", "a", "b", "b", "b", "b", "c", "c", "c", "d"], top 5 frequent words are [“b”, “c”, "a", "d"]

import heapq
class Solution(object):
  def topKFrequent(self, combo, k):
    dic = {}
    heap = []
    res = []
    for s in combo:
      if s in dic:
        dic[s] += 1
      else:
        dic[s] = 1
    for s in dic:
      heapq.heappush(heap,(-dic[s],s))
    if k <= len(heap):
      for i in xrange(k):
        res.append(heapq.heappop(heap)[1])
      return res
    else:
      for i in xrange(len(heap)):
        res.append(heapq.heappop(heap)[1])
      return res