// Top K Frequent Elements - https://leetcode.com/problems/top-k-frequent-elements/description/?source=submission-ac

class Solution {

    // Time Complexity: O(nlogn)
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();

        for (int i = 0; i < nums.length; i++) {
            int count = map.getOrDefault(nums[i], 0);
            map.put(nums[i], count + 1);
        }

        // Sorted in descending order by count
        PriorityQueue<Element> pq = new PriorityQueue<>(
            (el1, el2) -> Integer.compare(el2.j, el1.j)
        );

        for(Map.Entry<Integer, Integer> entry: map.entrySet()) {
            System.out.println(entry.getKey() + " " + entry.getValue());
            pq.offer(new Element(entry.getKey(), entry.getValue()));
        }

        int[] result = new int[k];
        for (int i = 0; i < k && !pq.isEmpty(); i++) {
            Element el = pq.poll();
            result[i] = el.i;
        }

        return result;
    }

    class Element {
        int i, j;

        Element(int i, int j) {
            this.i = i;
            this.j = j;
        }
    }

    Time Complexity: O(Nlogk)
    Space Complexity: O(N + k)
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> count = new HashMap<>();

        for (int n: nums) {
            count.put(n, count.getOrDefault(n, 0) + 1);
        }

        // Sorted in descending order by count
        Queue<Integer> minHeap = new PriorityQueue<>(
            (n1, n2) -> 
                Integer.compare(count.get(n1), 
                    count.get(n2))
        );

        for(int n: count.keySet()) {
            minHeap.offer(n);
            if(minHeap.size() > k) {
                minHeap.poll();
            }
        }

        int[] top = new int[k];
        for (int i = k - 1; i >= 0; i--) {
            top[i] = minHeap.poll();
        }

        return top;
    }
}