Kth Largest and Smallest Elements Using Heap

Introduction to Heaps

A heap is a special tree-based data structure that satisfies the heap property. There are two types of heaps:

1. Max Heap: The value of each node is greater than or equal to the values of its children. The largest element is at the root.
2. Min Heap: The value of each node is less than or equal to the values of its children. The smallest element is at the root.

Heaps are commonly used to implement priority queues, making them a versatile choice for finding the Kth largest or Kth smallest elements.

Problem Definition

Given an array of integers and an integer K, our goal is to find:

• The Kth largest element in the array.
• The Kth smallest element in the array.

Understanding the Approach

To achieve this using heaps, we can utilize the following strategies:

• For finding the Kth largest element, we can maintain a min heap of size K. By pushing all elements of the array into this min heap, we ensure that once we reach K elements, the root of the heap represents the Kth largest element.

• For finding the Kth smallest element, we can employ a max heap of size K. In this case, the root will represent the Kth smallest element after processing the array.

Finding the Kth Largest Element

Let's look at the implementation for finding the Kth largest element.

import java.util.PriorityQueue;

public class KthLargest {
    public static int findKthLargest(int[] nums, int k) {
        // Create a min heap with the default constructor
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();

        // Add elements to the min heap
        for (int num : nums) {
            minHeap.add(num);
            // Ensure the heap size does not exceed K
            if (minHeap.size() > k) {
                minHeap.poll(); // Removes the smallest (root) element
            }
        }
        return minHeap.peek(); // The root now is the Kth largest
    }

    public static void main(String[] args) {
        int[] nums = {3, 2, 1, 5, 6, 4};
        int k = 2; // Looking for the 2nd largest element
        System.out.println("The " + k + "th largest element is: " + findKthLargest(nums, k));
    }
}

Explanation:

1. We use a PriorityQueue to create a min heap.
2. We iterate over each element in the array.
3. We add each element to the min heap and ensure that its size does not exceed K by removing the smallest element when necessary.
4. After processing all elements, the smallest element in the min heap will be the Kth largest element.

Finding the Kth Smallest Element

Now, let’s implement the solution for finding the Kth smallest element.

import java.util.Collections;
import java.util.PriorityQueue;

public class KthSmallest {
    public static int findKthSmallest(int[] nums, int k) {
        // Create a max heap using a custom comparator
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

        // Add elements to the max heap
        for (int num : nums) {
            maxHeap.add(num);
            // Ensure the heap size does not exceed K
            if (maxHeap.size() > k) {
                maxHeap.poll(); // Removes the largest (root) element
            }
        }
        return maxHeap.peek(); // The root now is the Kth smallest
    }

    public static void main(String[] args) {
        int[] nums = {3, 2, 1, 5, 6, 4};
        int k = 2; // Looking for the 2nd smallest element
        System.out.println("The " + k + "th smallest element is: " + findKthSmallest(nums, k));
    }
}

Explanation:

1. We create a max heap using the Collections.reverseOrder() comparator.
2. For each element in the input array, we add it to the max heap while ensuring that its size does not exceed K by removing the largest element as necessary.
3. The root of the max heap after processing all elements will represent the Kth smallest element.

Complexity Analysis

1. Time Complexity: Both approaches have a time complexity of O(N log K), where N is the number of elements in the array. This is because each insertion into the heap takes O(log K) time, and we perform N insertions.

2. Space Complexity: The space complexity is O(K) due to the storage required for keeping K elements in the heap.

Conclusion

Using heaps to find the Kth largest and smallest elements in an array provides an efficient method that boils down the problem to managing a fixed-size heap. Whether using a min heap or a max heap, we can navigate through the array elements seamlessly and arrive at our required results promptly. Implementing these solutions in Java provides a practical insight into heap operations and enhances your understanding of advanced data structures.