Lintcode 471. Top K Frequent Words

Given a list of words and an integer k, return the top k frequent words in the list.

Example

Example 1:

Input:
  [
    "yes", "lint", "code",
    "yes", "code", "baby",
    "you", "baby", "chrome",
    "safari", "lint", "code",
    "body", "lint", "code"
  ]
  k = 3
Output: ["code", "lint", "baby"]

Example 2:

Input:
  [
    "yes", "lint", "code",
    "yes", "code", "baby",
    "you", "baby", "chrome",
    "safari", "lint", "code",
    "body", "lint", "code"
  ]
  k = 4
Output: ["code", "lint", "baby", "yes"]

Challenge

Do it in O(nlogk) time and O(n) extra space.

Notice

You should order the words by the frequency of them in the return list, the most frequent one comes first. If two words has the same frequency, the one with lower alphabetical order come first.

Code (Java):

public class Solution {
    /**
     * @param words: an array of string
     * @param k: An integer
     * @return: an array of string
     */
    public String[] topKFrequentWords(String[] words, int k) {
        // write your code here
        String[] ans = new String[k];
        if (words == null || words.length == 0) {
            return ans;
        }
        
        // step 1: do the freq count
        Map<String, Integer> wordToCountMap = new HashMap<>();
        for (String word : words) {
            int count = wordToCountMap.getOrDefault(word, 0);
            count += 1;
            wordToCountMap.put(word, count);
        }
        
        // step 2: use a min-heap to get the top k 
        PriorityQueue<Pair> minHeap = new PriorityQueue<>(new MyPairComparator());
        
        for (String word : wordToCountMap.keySet()) {
            int count = wordToCountMap.get(word);
            minHeap.offer(new Pair(word, count));
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }
        
        // step 3: poll the top k from the min heap
        int i = 0;
        while (!minHeap.isEmpty()) {
            ans[i++] = minHeap.poll().word;
        }
        
        // step 4: reverse the ans 
        Collections.reverse(Arrays.asList(ans)); 
        
        return ans;
    }
}

class Pair {
    String word;
    int count;
    
    public Pair(String word, int count) {
        this.word = word;
        this.count = count;
    }
}

class MyPairComparator implements Comparator<Pair> {
    @Override
    public int compare(Pair a, Pair b) {
        if (a.count - b.count != 0) {
            return a.count - b.count;
        }
        
        return b.word.compareTo(a.word);
    }
}