k самых частых элементов

Источник: https://leetcode.com/problems/top-k-frequent-elements/

Дан массив целых чисел nums и число k, вернуть k самых частых элементов. Самые частые элементы можно вернуть в любом порядке.

Пример 1:

Дано: nums = [1,1,1,2,2,3], k = 2
Результат: [1,2]

Пример 2:

Дано: nums = [1], k = 1
Результат: [1]

Ограниечения:

• 1 <= nums.length <= 105
• k в интервале [1, число уникальных жлементов в массиве].
• Изветно, что решение существует и уникально.

Алгоритм должен быть быстрее O(n log n), где n - размер массива (то есть обойтись без сортировки).

 

Имплементация

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        // O(1) time
        if (k == nums.length) {
            return nums;
        }
        
        // 1. строим hash map : частота элементов
        // O(N)
        Map<Integer, Integer> count = new HashMap();
        for (int n: nums) {
          count.put(n, count.getOrDefault(n, 0) + 1);
        }

        // новая КУЧА 'менее частый элемент в вершине'
        Queue<Integer> heap = new PriorityQueue<>(
            (n1, n2) -> count.get(n1) - count.get(n2));

        // 2. держим топ-k элементов (по частоте)в куче
        // O(N log k) < O(N log N), поскольку k<N
        for (int n: count.keySet()) {
          heap.add(n); // добавляем в кучу уникальный элемент
// он за log k попадает или в вершину, или оседает в куче
          if (heap.size() > k) heap.poll(); // держим в куче k элементов

// если их становится больше, то удаляем элемент из кучи, то есть из вершины,
// то есть менее частый, а это занимает log k
        }

        // 3. выводим результат для k оставшихся (самых частых) элементов
        // O(k log k)
        int[] top = new int[k];
        for(int i = k - 1; i >= 0; --i) {
            top[i] = heap.poll();
        }
        return top;
    }
}