// Runtime: 8 ms, faster than 62.46% of Java online submissions for Top K Frequent Words.
// Memory Usage: 38.6 MB, less than 67.86% of Java online submissions for Top K Frequent Words.

class Solution {
    public List<String> topKFrequent(String[] words, int k) {
        Map<String, Pair> countMap = new TreeMap<>();
        for (String word : words) {
            Pair pair = countMap.get(word);
            if (pair == null) {
                pair = new Pair();
                pair.word = word;
            }
            pair.count++;
            countMap.put(word, pair);
        }

        List<Pair> values = new ArrayList<>(countMap.values());
        values.sort(new Comparator<Pair>() {
            @Override
            public int compare(Pair o1, Pair o2) {
                return o2.count - o1.count;
            }
        });

        List<Pair> pairs = values.subList(0, k);
        ArrayList<String> result = new ArrayList<>();
        for (Pair pair : pairs) {
            result.add(pair.word);
        }
        return result;
    }

    static class Pair {
        String word;
        int count;
    }
}