Kth Largest/Smallest Element in an Array problem and solution in Java and Python

“Kth Largest/Smallest Element in an Array” is a classic FAANG problem. The efficient solution uses a heap, though other approaches like QuickSelect are also common.

Let’s go through Java and Python solutions.

Problem Statement

• Input: nums = [3,2,1,5,6,4], k = 2
• Output (Kth largest): 5
• Output (Kth smallest): 2

1️⃣ Java – Using Heap

Kth Largest Element

import java.util.PriorityQueue;

public class KthLargestElement {
    public static int findKthLargest(int[] nums, int k) {
        // Min-heap of size k
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();

        for (int num : nums) {
            minHeap.add(num);
            if (minHeap.size() > k) {
                minHeap.poll(); // remove smallest
            }
        }

        return minHeap.peek(); // kth largest
    }

    public static void main(String[] args) {
        int[] nums = {3,2,1,5,6,4};
        int k = 2;
        System.out.println(findKthLargest(nums, k)); // Output: 5
    }
}

Kth Smallest Element

• Use a max-heap of size k instead:

PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
for (int num : nums) {
    maxHeap.add(num);
    if (maxHeap.size() > k) {
        maxHeap.poll(); // remove largest
    }
}
return maxHeap.peek();

2️⃣ Python – Using Heap

Kth Largest Element

import heapq

def find_kth_largest(nums, k):
    min_heap = []
    for num in nums:
        heapq.heappush(min_heap, num)
        if len(min_heap) > k:
            heapq.heappop(min_heap)
    return min_heap[0]

# Example
nums = [3,2,1,5,6,4]
k = 2
print(find_kth_largest(nums, k))  # Output: 5

Kth Smallest Element

import heapq

def find_kth_smallest(nums, k):
    max_heap = []
    for num in nums:
        heapq.heappush(max_heap, -num)  # store negative for max-heap
        if len(max_heap) > k:
            heapq.heappop(max_heap)
    return -max_heap[0]

print(find_kth_smallest(nums, k))  # Output: 2

3️⃣ Alternative – QuickSelect

• Average O(n) time, worst-case O(n²)
• Works by partitioning like QuickSort.
• Often asked as a follow-up in interviews for faster than O(n log k).

⚡ Key Points for FAANG

1. Heap approach: O(n log k), simple to implement.
2. QuickSelect approach: O(n) average, better for large arrays.
3. Clarify whether kth largest or smallest is requested.
4. For Python, remember heapq is min-heap by default; use negative numbers for max-heap.