Find Kth Largest Element in Array

Given an array of huge numbers of elements, find the Kth largest element. For example, in the given array below, the 3rd largest would be 10.

[7, 10, 4, 3, 20, 15, 2]
Output - 10

One simple approach for this would be sorting the complete array and return the element present at index (n - 2). In this case, if we apply merge sort the worst time complexity would be O(nlogn) but this could be more optimized to O((nlogk) as we can find our kth element by sorting only k elements. We are actually unnecessarily sorting (n - k) elements to find the Kth largest or smallest element.

So, whenever we want to find any Kth largest or smallest element, we can achieve this with a time complexity of O(nlogk) using heap sort.

Heapsort is an in-place algorithm, but it is not a stable sort.

Heap is a special case of balanced binary tree data structure. A heap can be a max heap or min heap.

In case of max heap the value of the root node is more than or equal to either of its children whereas in case of min heap the value of the root node is less than or equal to either of its children.

There are some thumb rule in order to find any Kth smallest or largest element using heap sort.

To find Kth largest element, we always need to construct a min-heap and the top element in the heap would be our Kth largest element.

To find Kth smallest element, we always need to construct a max-heap and the top element in the max-heap would be our Kth smallest element.

Min Heap Implementation

As we know, the top element of min-heap is less than or equal to either of its children. Hence, to implement this in Java, we can directly use the inbuilt collection in Java - PriorityQeue.

The elements of the priority queue are ordered according to their natural ordering. The top element has always the smallest element. Hence, we can directly use this collection for the min-heap implementation.

Kth Largest Element Implementation

The algorithm for this would be

• Iterate each element of the given array
• Inserrt the element in the min-heap.
• After each insertion, check the size of heap. If the heap size exceeds K element, remove the head element.
• Repeat above process for all tthe elementts of the given array.
• Now, the top element of the min-heap would be our Kth largest element.

Below is the implementation for the same.

public class MinHeap {

    public static int kthLargestElement(int k, int[] array){
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        int size = array.length;
        for (int i = 0; i < size; i++){
            minHeap.add(array[i]);
            if (minHeap.size() > k){
                minHeap.poll();
            }
        }
        return minHeap.peek();
    }

    public static void main(String[] args) {
        int[] array = {7, 10, 4, 3, 20, 15, 2};
        System.out.println(MinHeap.kthLargestElement(3, array));
    }
}

Kth Smallest Element

In a similar way, we can also find the Kth smallest element. To find the Kth smallest element, we require to construct a max-heap.

In case of max heap, the largest element would be at the top of the heap and hence the order of priority queue needs to be reversed.

public class MaxHeap {

    public static int  kthSmallestElement(int  k, int [] array){
        PriorityQueue<Integer> maxHeap = new  PriorityQueue<>(Collections.reverseOrder());
        int  length = array.length;
        for  (int  i = 0; i < length; i++){
            maxHeap.add(array[i]);
            if  (maxHeap.size() > k){
                maxHeap.poll();
            }
        }
        return  maxHeap.peek();
    }

    public static void  main(String[] args) {
        int [] array = {1, 3, 8, 9, 4, 7, 6};
        System.out .println(MaxHeap.kthSmallestElement(3, array));
    }
}

Conclusion

In this article, we implemented the algorithm to find the kth largest and smallest element using heap in Java.