Heaps: Mastering the Top-K Elements Pattern

A Heap is a specialized tree-based data structure that satisfies the heap property: in a Min-Heap, the parent is always smaller than its children, and in a Max-Heap, the parent is always larger. In Kotlin, we implement this using the PriorityQueue class.

Heaps are the go-to structure for “Top-K” problems because they allow us to maintain a limited set of the most (or least) significant elements without sorting the entire dataset.

---

1. Top K Frequent Elements

The goal here is to find the elements that appear most often in an array. This requires two steps:

1. Building a Frequency Map to count occurrences.
2. Using a Min-Heap of size $K$ to keep track of the highest frequencies.

fun topKElements(inputArray: Array<Int>, k: Int): List<Int> {
    val frequencyTable = mutableMapOf<Int, Int>()

    // Step 1: Count frequencies
    for (element in inputArray) {
        frequencyTable[element] = frequencyTable.getOrDefault(element, 0) + 1
    }

    // Step 2: Min-Heap based on frequency
    val pq = PriorityQueue<Pair<Int, Int>> { a, b ->
        if (a.second == b.second) a.first.compareTo(b.first) 
        else a.second.compareTo(b.second)
    }

    frequencyTable.forEach { (key, value) ->
        pq.add(Pair(key, value))
        // Maintain heap size K
        if (pq.size > k) {
            pq.poll()
        }
    }

    val op = mutableListOf<Int>()
    while (pq.isNotEmpty()) {
        op.add(pq.poll().first)
    }

    return op.reversed()
}

2. K-Largest and K-Smallest Elements

Finding the K largest or smallest elements in an unsorted array is a classic problem that can be solved in O(N \log K) time using a Heap.

K-Largest (Using Min-Heap)

To find the largest elements, we use a Min-Heap. Why? Because a Min-Heap keeps the smallest of the “large” elements at the top (peek). If a new element is larger than the top, we remove the top and add the new one.

fun firstKLargestElements(inputArray: Array<Int>, k: Int): List<Int> {
    if(k >= inputArray.size) return emptyList()

    val pq = PriorityQueue<Int>(k) // Default is Min-Heap
    for(i in 0 until k) pq.add(inputArray[i])

    for(i in k until inputArray.size) {
        if(pq.peek() < inputArray[i]) {
            pq.poll()
            pq.add(inputArray[i])
        }
    }

    return pq.toList().sortedDescending()
}

K-Smallest (Using Max-Heap)

To find the smallest elements, we use a Max-Heap (created by passing reverseOrder()). This keeps the largest of the “small” elements at the top.

fun firstKSmallestElements(inputArray: Array<Int>, k: Int): List<Int> {
    if(k >= inputArray.size) return emptyList()

    // Max-priority Queue
    val pq = PriorityQueue<Int>(k, reverseOrder())
    for(i in 0 until k) pq.add(inputArray[i])

    for(i in k until inputArray.size) {
        if(pq.peek() > inputArray[i]) {
            pq.poll()
            pq.add(inputArray[i])
        }
    }

    return pq.toList().sorted()
}

Why use a Heap instead of Sorting?

Method	Time Complexity	Space Complexity
Full Sorte	O(N \log N)	O(1) or O(N)
Heap (Top-K)	O(N \log K)	$O(K)$

When K is much smaller than N, the Heap approach is significantly faster and more memory-efficient, especially for streaming data where you don’t have the full array upfront.

Summary

• PriorityQueue is Kotlin’s implementation of a Heap.
• Use a Min-Heap for “Largest” problems.
• Use a Max-Heap for “Smallest” problems.
• The “Top-K” pattern is essential for ranking systems, frequency analysis, and scheduling tasks.

Source Code

🚀 Code: All the code for this tutorial is available in my DSA-Kotlin Repository.