【Lintcode】471. Top K Frequent Words

题目地址：

https://www.lintcode.com/problem/top-k-frequent-words/description

给定一个单词的数组，问其出现频率最高的$k$个单词是哪些。如果有出现频率相同的单词，则认为字典序小的那个入选。

思路是用小顶堆。因为我们需要选出现频率最高并且字典序小的单词，所以我们构造小顶堆的时候，就让出现频率低和字典序大的单词优先级更高。注意，这里要先对$k=0$的情况进行判断。代码如下：

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

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
        // 如果k = 0则返回空数组
        if (k == 0) {
            return new String[0];
        }
        // 先用一个哈希表统计每个单词的出现频率
        Map<String, Integer> freq = new HashMap<>();
        for (String word : words) {
            freq.put(word, freq.getOrDefault(word, 0) + 1);
        }
        // 构造一个小顶堆，单词出现频率小的优先，如果频率一样，则字典序大的优先
        PriorityQueue<String> minHeap = new PriorityQueue<>((s1, s2) -> {
            if (!freq.get(s1).equals(freq.get(s2))) {
                return freq.get(s1) < freq.get(s2) ? -1 : 1;
            } else {
                for (int i = 0, j = 0; i < s1.length() && j < s2.length(); i++, j++) {
                    if (s1.charAt(i) != s2.charAt(j)) {
                        return s1.charAt(i) > s2.charAt(j) ? -1 : 1;
                    }
                }
                
                return s1.length() > s2.length() ? -1 : 1;
            }
        });
        
        // 如果发现新来的单词频率比堆顶更高，或者虽然频率一样但字典序靠前，则将堆顶踢出堆，将新来的词加入堆
        for (String word : freq.keySet()) {
            if (minHeap.size() < k) {
                minHeap.offer(word);
            } else {
                if (freq.get(word) > freq.get(minHeap.peek()) ||
                        (freq.get(word).equals(freq.get(minHeap.peek())) && word.charAt(0) < minHeap.peek().charAt(0))) {
                    minHeap.poll();
                    minHeap.offer(word);
                }
            }
        }
        // 逆序对res赋值
        String[] res = new String[k];
        for (int i = k - 1; i >= 0; i--) {
            res[i] = minHeap.poll();
        }
        
        return res;
    }
}

时间复杂度$O(nl\log k)$，空间$O(nl)$，$n$为单词个数，$l$为最长单词长度。