【Lintcode】550. Top K Frequent Words II

题目地址：

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

要求设计一个数据结构，可以实现添加字符串和topk操作，并且有容量限制。topk返回的是出现次数最高的$k$个字符串，如果两个字符串出现次数相等，则字典序小者优先。

我们可以用一个HashMap记录每个字符串出现了多少次。这里由于要动态维护topk的字符串，用堆的话，即使字符串的出现次数的计数改变了，堆也不会自己调整，而要先删除后调整，而java在堆中删除的效率是很低的，$O(n)$复杂度，所以不适合用堆（当然如果手写哈希堆也可以）。而用TreeSet的话，删除可以达到$O(\log n)$，快很多。

但是用TreeSet尤其要注意几个问题：
1、对象声明的时候不能用Set<String> set = new TreeSet<>();而必须用TreeSet<String> set = new TreeSet<>();，这是因为在add的操作中，如果新来的字符串要替代set里已有的字符串，需要将set里优先级最低的字符串删掉，需要用到TreeSet的pollLast()方法，而Set里是没有这个方法的；
2、TreeSet同样会遇到，一个字符串计数改变但是不会自己调整的问题，这时候需要先将字符串从set里删掉，再加回去；但是，由于TreeSet的比较器会写成，出现次数多者优先，次数一样多则字典序小者优先，这样会导致一个问题：由于TreeSet是按照比较器判断元素相等的，并不是按照hashCode和equals比较的，所以在判断两个字符串哪个大哪个小的时候，需要保证map里存在这两个字符串，否则会NPE。所以要从set里添加或者判断元素是否存在的时候，一定要保证map里有这个元素；
3、如果set里已经有了某个字符串，而这个字符串又来了，一定要先将该字符串从set里删掉，再将其计数加一。否则的话，先计数加一再删除，这个字符串就再也找不到了，因为计数不同了。先删除，再计数加一，再加回去，这样才能保证set里的字符串的顺序是维护正确的。

考虑到以上几点后，我们将add的逻辑叙述如下：
1、如果新来的字符串s已经存在于set里，那我们要做的事是维护其在set里的顺序正确，也就是要先从set里删掉s，增加计数，再加回来。但是，一旦s不在map里，我们调set.contains(s)或set.remove()的时候就会导致NPE。所以我们可以先看一下map里有没有s，再判断set里有没有s，这样可以避免NPE，并且由于计数没变，所以相同的字符串一定能被找到。如果map里有的话，再看set，判断在里面的话，就将其删除，再计数加一，再加回来；如果判断不在里面，那就计数加一之后，再把这个字符串加进去即可。
2、如果新来的字符串s不存在于map里，那就更不可能在set里，所以可以直接计数加一然后加进set；
3、如果set的size大于k了，就把优先级最低的那个字符串删掉。

按照上面的操作顺序就可以万无一失。具体请看代码。代码如下：

import java.util.*;

public class TopK {
    
    private int k;
    private Map<String, Integer> map;
    // 这里要用TreeSet，不能写Set
    private TreeSet<String> set;
    
    /*
     * @param k: An integer
     */
    public TopK(int k) {
        // do intialization if necessary
        this.k = k;
        map = new HashMap<>();
        // 比较器是按照出现次数多者优先，次数一样按字典序小的优先
        set = new TreeSet<>((s1, s2) -> {            
            if (!map.get(s1).equals(map.get(s2))) {
                return -Integer.compare(map.get(s1), map.get(s2));
            }
            
            return s1.compareTo(s2);
        });
    }
    
    /*
     * @param word: A string
     * @return: nothing
     */
    public void add(String word) {
        // write your code here
        // 先看在不在map里。如果在，再看是否在set里，
        // 如果在，就删掉后计数加一，否则直接计数加一；之后再加进set。
        // 如果不在map里，那就计数加一，之后加进set
        if (map.containsKey(word)) {
            if (set.contains(word)) {
                set.remove(word);
            }
            
            map.put(word, map.get(word) + 1);
        } else {
            map.put(word, 1);
        }
    
        set.add(word);
        // size大于k了就把最小的那个删掉
        if (set.size() > k) {
            set.pollLast();
        }
    }
    
    /*
     * @return: the current top k frequent words.
     */
    public List<String> topk() {
        // write your code here
        return new ArrayList<>(set);
    }
}

add时间复杂度$O(\log k)$，空间$O(n)$。