Skip to main content

Heaps

Kth Largest Element in an Array Solution In C++/Python/Java/JS

Problem Description

Given an integer array nums and an integer k, return the kth largest element in the array.
Note that it is the kth largest element in the sorted order, not the kth distinct element.
Can you solve it without sorting?

Example

Input: nums = [3,2,1,5,6,4], k = 2
Output: 5


Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4

Constraints

  • 1 <= k <= nums.length <= 10^5
  • -10^4 <= nums[i] <= 10^4
💡
Try it Yourself First
Finding the k-th largest element in an array is a well-known problem in computer science, and several efficient techniques have been developed to solve it. Let's explore different approaches to tackle this problem effectively

Brute Force Approach

Intuition

The goal is to find the Kth largest element in the array. The most straightforward way to achieve this is by sorting the array in descending order. This ensures that the largest elements appear first, making it easy to directly access the Kth largest element at index K−1 (zero-based indexing).

Since sorting arranges all elements in the required order, this approach guarantees correctness. However, it takes O(NlogN) time due to the sorting step, which may not be optimal for very large inputs.

Approach

Step 1: Sort the Array in Descending Order

  • Sort the given array in descending order to arrange elements from largest to smallest.

Step 2: Identify the Kth Largest Element

  • Since the array is now sorted, the Kth largest element will be at index K−1 (zero-based indexing).

Step 3: Return the Result

  • Retrieve and return the element at index K−1 as the final answer.

Dry Run

Input:

  • nums = [3, 2, 1, 5, 6, 4], k = 2

Step 1: Sort the Array in Descending Order

  • Sorting the array in descending order:
  • nums = [6, 5, 4, 3, 2, 1]

Step 2: Identify the Kth Largest Element

  • Since we need the 2nd largest element, we pick the element at index k - 1 = 2 - 1 = 1.

Step 3: Return the Result

  • The element at index 1 is 5, so the output is:
  • Output: 5

Final max_elegance = 5

Code for All Languages
C++ 
// Approach: Sorting-based Solution to Find Kth Largest Element
// Language: C++

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        // Step 1: Sort the array in descending order
        sort(nums.begin(), nums.end(), greater<int>());

        // Step 2: Return the Kth largest element (0-based index K-1)
        return nums[k - 1];
    }
};

Java 
// Approach: Sorting-based Solution to Find Kth Largest Element
// Language: Java

import java.util.Arrays;

class Solution { 
    public int findKthLargest(int[] nums, int k) {
        // Step 1: Sort the array in descending order
        Arrays.sort(nums);
        
        // Step 2: Return the Kth largest element (0-based index nums.length - k)
        return nums[nums.length - k];
    }
}

Python 
# Approach: Sorting-based Solution to Find Kth Largest Element
# Language: Python

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        # Step 1: Sort the array in descending order
        nums.sort(reverse=True)

        # Step 2: Return the Kth largest element (0-based index k-1)
        return nums[k - 1]

Javascript 
// Approach: Sorting-based Solution to Find Kth Largest Element
// Language: JavaScript

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var findKthLargest = function(nums, k) {
    // Step 1: Sort the array in descending order
    nums.sort((a, b) => b - a);

    // Step 2: Return the Kth largest element (0-based index k-1)
    return nums[k - 1];
};

Time Complexity Analysis: O(n log n)

Let n be the number of items in the input array.

Sorting Step:

  • We sort the array in descending order, which takes O(nlogn) time.

Accessing the Kth Largest Element:

  • After sorting, we directly access the element at index k−1, which takes O(1) time.

Final Time Complexity:

  • The dominant factor is the sorting step, so the overall time complexity is O(nlogn)

Final Time Complexity: O(n log n)

Space Complexity Analysis: O(1)

Auxiliary Space Complexity: O(1)

  • Sorting is done in place, so no extra space is required.
  • We only use a few integer variables to store indices and the result, which takes O(1) space.

Total Auxiliary Space Complexity: O(1)

Total Space Complexity: O(1)

  • The input array nums is already given, and we do not use any additional data structures that scale with n.
  • Since we do not use extra arrays, stacks, or hash sets, the total space complexity remains O(1).

Final Space Complexity:O(1).

The brute-force approach involves sorting the array and then selecting the k-th largest element, which is inefficient due to its O(n log n) complexity. Using a min-heap, we dynamically maintain the k largest elements, ensuring efficient tracking while removing smaller elements. This allows us to find the k-th largest element in O(n log k) time, significantly improving performance over sorting.

Optimal Approach

Intuition

To optimize this, a Min-Heap (Priority Queue) can be used instead of sorting. The idea is to maintain a heap of size k while iterating through the array, ensuring that it always contains the top k largest elements encountered so far. We begin by inserting the first k elements into the heap. For each subsequent element, we compare it with the smallest element in the heap (the root). If the current element is larger, we remove the root and insert the new element, keeping only the k largest elements in the heap. By the end of the iteration, the root of the heap represents the kth largest element in the array.

Approach

Step 1: Create a Min-Heap

  • A min-heap (priority queue) is used to store the k largest elements encountered so far.
  • This ensures that the smallest element among them is always at the top.

Step 2: Insert the First k Elements into the Heap

  • Iterate through the first k elements of the array.
  • Insert each element into the min-heap.
  • The heap maintains the smallest element at the top.

Step 3: Process the Remaining Elements

  • Iterate through the remaining elements in the array.
  • If an element is greater than the smallest element in the heap, replace it.
  • Remove the smallest element from the heap.
  • Insert the new element into the heap to maintain the k largest elements.

Step 4: Extract the k-th Largest Element

  • The top element of the min-heap now represents the k-th largest element in the array.
  • Return this element as the final result.

Dry Run

Input: nums = [3,2,1,5,6,4], k = 2

Step 1: Create a Min-Heap

  • Initialize a min-heap to store the k largest elements encountered.

Step 2: Insert the First k = 2 Elements into the Heap

  • Pick 3 → Heap = [3]
  • Pick 2 → Heap = [2, 3]

Step 3: Process the Remaining Elements

  • Next element 1 → It is smaller than the smallest element in the heap (2), so it is ignored. Heap remains [2, 3].
  • Next element 5 → It is greater than 2, so replace 2 with 5. Heap = [3, 5].
  • Next element 6 → It is greater than 3, so replace 3 with 6. Heap = [5, 6].
  • Next element 4 → It is smaller than 5, so it is ignored. Heap remains [5, 6].

Step 4: Extract the k-th Largest Element

  • The top element in the min-heap is 5, which is the 2nd largest element in the array.

Final Output: 5

Code for All Languages
C++ 
// Approach: Heap-Based Solution to Find the K-th Largest Element
// Language: C++

class Solution {
public:
    int findKthLargest(std::vector<int>& nums, int k) 
    {
        // Step 1: Create a min-heap (priority queue).
        priority_queue<int, vector<int>,greater<int>> minHeap;

        // Step 2: Insert the first k elements into the heap.
        for (int i = 0; i < k; i++) 
        {
            minHeap.push(nums[i]);
        }

        // Step 3: Process the remaining elements.
        for (int i = k; i < nums.size(); i++) 
        {
            // If the current element is larger than the smallest in the heap, replace it.
            if (nums[i] > minHeap.top()) 
            {
                minHeap.pop();
                minHeap.push(nums[i]);
            }
        }

        // Step 4: Return the k-th largest element, which is at the top of the heap.
        return minHeap.top();
    }
};

Java 
// Approach: Heap-Based Solution to Find the K-th Largest Element
// Language: Java

import java.util.PriorityQueue;

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // Step 1: Create a min-heap (priority queue).
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();

        // Step 2: Insert the first k elements into the heap.
        for (int i = 0; i < k; i++) {
            minHeap.add(nums[i]);
        }

        // Step 3: Process the remaining elements.
        for (int i = k; i < nums.length; i++) {
            // If the current element is larger than the smallest in the heap, replace it.
            if (nums[i] > minHeap.peek()) {
                minHeap.poll();
                minHeap.add(nums[i]);
            }
        }

        // Step 4: Return the k-th largest element, which is at the top of the heap.
        return minHeap.peek();
    }
}

Python 
# Approach: Heap-Based Solution to Find the K-th Largest Element
# Language: Python

import heapq

class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        # Step 1: Create a min-heap (priority queue).
        min_heap = []

        # Step 2: Insert the first k elements into the heap.
        for num in nums[:k]:
            heapq.heappush(min_heap, num)

        # Step 3: Process the remaining elements.
        for num in nums[k:]:
            # If the current element is larger than the smallest in the heap, replace it.
            if num > min_heap[0]:
                heapq.heappushpop(min_heap, num)

        # Step 4: Return the k-th largest element, which is at the top of the heap.
        return min_heap[0]

Javascript 
// Approach: Heap-Based Solution to Find the K-th Largest Element
// Language: JavaScript

class MinHeap {
    constructor() {
        // Initialize an empty heap array
        this.heap = [];
    }

    push(val) {
        // Add the value to the heap
        this.heap.push(val);

        // Sort the heap to maintain the min-heap property
        this.heap.sort((a, b) => a - b);
    }

    pop() {
        // Remove and return the smallest element from the heap
        return this.heap.shift();
    }

    top() {
        // Return the smallest element in the heap without removing it
        return this.heap[0];
    }

    size() {
        // Return the current size of the heap
        return this.heap.length;
    }
}

var findKthLargest = function(nums, k) {
    // Step 1: Create a min-heap (priority queue)
    let minHeap = new MinHeap();

    // Step 2: Insert the first k elements into the heap
    for (let i = 0; i < k; i++) {
        minHeap.push(nums[i]);
    }

    // Step 3: Process the remaining elements in the array
    for (let i = k; i < nums.length; i++) {
        // If the current number is greater than the smallest in the heap
        if (nums[i] > minHeap.top()) {
            // Remove the smallest element from the heap
            minHeap.pop();

            // Insert the current number into the heap
            minHeap.push(nums[i]);
        }
    }

    // Step 4: Return the k-th largest element, which is at the top of the heap
    return minHeap.top();
};

Time Complexity Analysis: O(n log k)

Let n be the number of elements in the input array and k be the position of the k-th largest element.

We use a Min-Heap of size k to keep track of the k largest elements efficiently.

Processing the first k elements: O(k log k)

  • We insert the first k elements into the heap.
  • Each insertion takes O(log k), so for k elements, this step takes O(k log k).

Processing the remaining (n - k) elements: O((n - k) log k)

  • For each of the remaining (n - k) elements:
  • If the element is larger than the heap's root, we remove the smallest element (O(log k)) and insert the new element (O(log k)).
  • Since this happens for (n - k) elements, it contributes O((n - k) log k) time.

Final Complexity: O(n log k)

Since k ≤ n, in the worst case, log k ≈ log n, so O(n log k) remains optimal compared to O(n log n) sorting.

Space Complexity Analysis: O(k) ⊆ O(n)

Heap Storage: O(k)

  • We maintain a Min-Heap of size k to store the k largest elements.
  • Since the heap size is at most k, it takes O(k) space, which is at most O(n) when k = n in the worst case.

Input Array Storage: O(n)

  • The input array itself takes O(n) space but is not considered extra space since it is part of the input.

Auxiliary Space: O(1)

  • A few integer variables (such as loop counters and temporary storage) take O(1) space.

Final Space Complexity: O(k) ⊆ O(n)

  • The dominant space usage comes from the hash map (O(n)) and min-heap (O(k) ⊆ O(n)).
  • Thus, the overall space complexity remains O(n).

Similar Problems

Now we have successfully tackled this problem, let’s try these similar problems.

Kth Largest Element in an Array

Given an integer array nums and an integer k, return the kth largest element in the array.Note that it is the kth largest element in the sorted order, not the kth distinct element.

Find Median from Data Stream

This problem requires continuously finding the median of a growing data stream. It uses two heaps (a max-heap for the left half and a min-heap for the right half) to efficiently balance and retrieve the median in O(log n) insertion and O(1) retrieval time, similar to how we maintain the k-th largest element in a min-heap.