NeetCode 150: Top K Frequent Elements

Table of contents

The Problem

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

The Approach

The goal is to first sort the elements by their frequency in nums & pop the first k indices.

By inserting the elements into a Map, where the element is the key & its frequency is the value, we assign each element its respective frequency.

We can then place these elements in a PriorityQueue, where each element’s priority is determined by its frequency. Using the poll() method of the PriorityQueue, we can remove the k most frequent elements.

My Solution

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        // Use a HashMap for fast insertion.
        Map<Integer, Integer> numsFrequency = new HashMap<Integer, Integer>();
        
        // Inserts the Integer & its frequency into a Map.
        for (int i = 0; i < nums.length; i++) {
            Integer currentInt = nums[i];

            if (!(numsFrequency.containsKey(currentInt))) {
                numsFrequency.put(currentInt, 1);
            } else {
                int currentFrequency = numsFrequency.get(currentInt);
                numsFrequency.replace(currentInt, currentFrequency + 1);
            }
        }
        // Inserts the elements of the Map into a PQ, where the order is
        // determined by the frequency of the element.
        Queue<Integer> numsPQ = new PriorityQueue<Integer>(
            // "b" compared to "a" for a descending PQ.
            (a, b) -> (numsFrequency.get(b)).compareTo(numsFrequency.get(a))
        );

        for (Integer current : numsFrequency.keySet()) {
            numsPQ.add(current);
        }

        // Removes the first k values from the PQ.
        int[] result = new int[k];
        for (int j = 0; j < k; j++) {
            result[j] = numsPQ.poll();
        }
        return result;
    }
}

Firstly, we initialize numsFrequency which is a HashMap to create key-value pairs of nums’s elements & their respective frequencies.

We then initialize numsPQ which is a PriorityQueue, passing a lamda expression which acts as a comparator that orders the elements by their frequency. The keys from numsFrequency are inserted into numsPQ.

The first k elements from numsPQ are removed & placed in an int[] which is returned.