How to Find the Kth Largest Element in an Array

The most obvious way to find the Kth largest element in an array would be to sort the array in descending order and then access its element at the index k - 1. However, the shortcoming of this approach comes from the fact that it doesn’t solve this problem in the most optimal way with regards to time complexity.

Brute-Force Solution (Time Complexity: O(nlogn), Space Complexity: O(1))

int Kth_largest(const std::vector<int>& array, int k)
{
   std::sort(array.rbegin(), array.rend());
  
   return array[k - 1];
}

To advance towards a more optimal solution we need to make use of another data structure that provides us with better time complexity guarantees. A priority queue is an excellent tool for this kind of improvement as it allows us to compare at most K elements instead of the whole array whenever we are sorting as we traverse through the array.

Min-Heap Priority Queue Solution (Time Complexity: O(nlogk), Space Complexity: O(k))

int Kth_largest(const std::vector<int>& array, int k)
{
   std::priority_queue<int, std::vector<int>, std::greater<int>> queue;
  
   for(auto val : array)
   {
      queue.push(val);
     
      if(queue.size() > k)
      {
          queue.pop();
      }
   }
  
    return queue.top();
}

This solution is a little more verbose than the first one, but it is still very simple and useful when time efficiency is of concern. Moreover, the previous implementation is also adaptable to the Kth smallest element problem by replacing the min-heap queue with a max-heap queue. Fortunately, the C++ STL provides an in-built function called nth_element that is implemented with partial sorting using introspective selection to offer an excellent linear average time complexity with the only downside of the rare case of a bad partition that leads to a worst possible time complexity of O(nlogn).

Nth Element Solution (Avg Time Complexity: O(n), Worst Time Complexity: O(nlogn), Space Complexity: O(1))

int Kth_largest(const std::vector<int>& array, int k)
{
   std::nth_element(array.begin(), array.begin() + k - 1, array.end(), std::greater<int>());
  
   return array[k - 1];
}

Conclusion

It is necessary to divert from the standard sorting of an array to find an optimal time efficient solution to the Kth largest element problem. Priority queues and more specialized partial sorting algorithms provide an excellent alternative to solving this issue.

Source Code