Kth Largest Element in an Array

Patrick Leaton
Problem Description

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

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

Example 2:

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

Note:
You may assume k is always valid, 1 ≤ k ≤ array's length.

 

Description taken from https://leetcode.com/problems/kth-largest-element-in-an-array/.

 

Problem Solution

#O(N(Log(N))) Time, O(N) Space
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = []
        for num in nums:
            heapq.heappush(heap, num)
        while len(heap) > k:
            heapq.heappop(heap)
        return heapq.heappop(heap)

Problem Explanation


The solution that may come to mind right away is to sort the array and return the kth from the last index.  That would be O(N(Log(N))) time and constant space.  

However, we can improve the speed of our solution to O(N(Log K)) time, but using linear space.  

This solution utilizes Python's min-heap, similar to the Kth Largest Element in a Stream problem.  We can push each element into the heap and pop off elements until the heap is size k, in which case our top element will be the kth largest.


We will start by initializing the heap.

        heap = []

 

Next, we will push all of the values from our input array onto the heap.

        for num in nums:
            heapq.heappush(heap, num)

 

We will then use a loop that will run while the heap size is greater than our k value.

This will pop the k-1 smallest elements so by the time the loop breaks, we will have the kth largest element as the minimum value in our heap.

        while len(heap) > k:
            heapq.heappop(heap)

 

We can pop off this top element and return it.

        return heapq.heappop(heap)


 

 

Additional Notes

This problem and other similar heap solutions can also be solved by using Quick Select.  This gives a better average time complexity but a worse worst-case time complexity.