Kth Largest Element Using Max Heap in DSA Go

Mental Model

We use a max heap to keep the largest numbers at the top so we can quickly find the kth largest by removing the largest elements one by one.

Analogy: Imagine a pile of books stacked by size with the biggest on top. To find the kth biggest book, you remove the biggest books from the top until you reach the kth one.

Max Heap Array Representation:
[ 50, 30, 40, 10, 20 ]
Heap Tree:
      50
     /  \
   30    40
  /  \
10   20
↑ root (largest)

Dry Run Walkthrough

Input: array: [20, 10, 40, 30, 50], k=3

Goal: Find the 3rd largest element using a max heap

Step 1: Build max heap from array

[50, 30, 40, 10, 20]

Why: We rearrange the array so the largest element is at the root for easy access

Step 2: Remove the largest element (50) from heap

[40, 30, 20, 10]

Why: Removing the largest element moves us closer to the kth largest

Step 3: Remove the next largest element (40) from heap

[30, 10, 20]

Why: We continue removing largest elements to reach the kth

Step 4: Remove the next largest element (30) from heap

[20, 10]

Why: The 3rd removed element is the 3rd largest in the original array

Result:

Heap after 3 removals: [20, 10]
3rd largest element: 30

Annotated Code

DSA Go

package main

import (
	"container/heap"
	"fmt"
)

// MaxHeap implements heap.Interface for a max heap
type MaxHeap []int

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] } // max heap
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *MaxHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *MaxHeap) Pop() interface{} {
	n := len(*h)
	val := (*h)[n-1]
	*h = (*h)[:n-1]
	return val
}

func findKthLargest(nums []int, k int) int {
	h := MaxHeap(nums)
	heap.Init(&h) // build max heap

	var kthLargest int
	for i := 0; i < k; i++ {
		kthLargest = heap.Pop(&h).(int) // remove largest k times
	}
	return kthLargest
}

func main() {
	nums := []int{20, 10, 40, 30, 50}
	k := 3
	result := findKthLargest(nums, k)
	fmt.Printf("%dth largest element: %d\n", k, result)
}

heap.Init(&h) // build max heap

convert array into max heap so largest element is at root

kthLargest = heap.Pop(&h).(int) // remove largest k times

remove the root (largest) element k times to get kth largest

Output

3th largest element: 30

Complexity Analysis

Time: O(n + k log n) because building the heap takes O(n) and each of the k removals takes O(log n)

Space: O(n) because the heap stores all n elements

vs Alternative: Compared to sorting the array O(n log n), this approach is faster when k is small because it avoids full sorting

Edge Cases

k is 1 (find largest element)

Returns the maximum element correctly

DSA Go

for i := 0; i < k; i++ {

k equals length of array (find smallest element)

Returns the smallest element after removing all larger ones

DSA Go

for i := 0; i < k; i++ {

array with duplicate elements

Duplicates are handled normally as heap stores all elements

DSA Go

heap.Init(&h)

When to Use This Pattern

When you need the kth largest or smallest element efficiently, use a heap to avoid full sorting and get the element by repeated removal.

Common Mistakes

Mistake: Using a min heap instead of a max heap without adjusting logic
Fix: Use a max heap or invert comparison so largest elements are on top

Mistake: Not building the heap before popping elements
Fix: Call heap.Init to build the heap before popping

Summary

Finds the kth largest element by building a max heap and removing the largest elements k times.

Use when you want the kth largest element without sorting the entire array.

The key insight is that a max heap keeps the largest element at the root for quick access.