Top K Frequent Words

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

Example:

Given

[

"yes", "lint", "code",

"yes", "code", "baby",

"you", "baby", "chrome",

"safari", "lint", "code",

"body", "lint", "code"

]

for k = 3, return ["code", "lint", "baby"].

for k = 4, return ["code", "lint", "baby", "yes"].

Challenge:

Do it in O(n) time with O(k) extra space approximation algorithms.

Solution:

This is a similar to Top k Largest Numbers II when we add all the words into a HashMap with word and frequency as key, value pairs. A k-length min PriorityQueue is used to track the largest k words.

Code:

class Pair {
    String s;
    int f;
    public Pair(String inS, int inF) {
        this.s = inS;
        this.f = inF;
    }
}

public class Solution {
    private Comparator<Pair> pairComparator = new Comparator<Pair>() {
        public int compare(Pair a, Pair b) {
            if (b.f == a.f) {
                return b.s.compareTo(a.s);
            }
            return a.f - b.f;
        }
    };
    public String[] topKFrequentWords(String[] words, int k) {
        if (words == null || words.length == 0) {
            return words;
        }
        if (k == 0) {
            return new String[] {};
        }

        Map<String,Integer> map = new TreeMap<String,Integer>();
        for (String s : words) {
            if (map.containsKey(s)) {
                map.put(s, map.get(s) + 1);
            } else {
                map.put(s, 1);
            }
        }
        //create a min queue stores only k elements
        Queue<Pair> q = new PriorityQueue(k, pairComparator);
        //add the new word in if its frequency is more than the minimal one in the min queue
        for (Map.Entry<String, Integer> entry : map.entrySet()) {
            Pair pair = new Pair(entry.getKey(), entry.getValue());
            if (q.size() < k) {
                q.offer(pair);
            } else if (pairComparator.compare(pair, q.peek()) > 0) {
                q.poll();
                q.offer(pair);
            }
        }
        String[] result = new String[k];
        for (int i = 0; i < k; i++) {
            result[i] = q.poll().s;
        }
        // reverse
        for (int i = 0; i < k / 2; i++) {
            String temp = result[i];
            result[i] = result[k - i - 1];
            result[k - i - 1] = temp;
        }
        return result;
    }
}