Heap Interview Problem: Merge $K$ Sorted Linked Lists & Merge $K$ Sorted Arrays

This comprehensive chapter explores two classical, high-impact algorithmic problems: Merge $K$ Sorted Linked Lists and Merge $K$ Sorted Arrays. By solving these problems, we transition our understanding of Heaps and Priority Queues to an advanced level, mastering custom comparators, space-time optimization, and specialized divide-and-conquer strategies.

1. Why Do We Need Optimized Approaches for $K$-Way Merging?

When dealing with $K$ distinct sorted data structures (whether Linked Lists or Arrays), the status quo naive approach involves flattening everything or repeatedly merging pairs.

• The Status Quo Problem: If you have $K$ sorted lists, each of size $N$, a brute-force combination yields a total of $N \times K$ elements.
• The Performance Bottleneck: Repeatedly appending or naively sorting unoptimized collections results in redundant comparisons ($O(N^2 K^2)$ or sluggish sorting overhead). We need a data structure that can dynamically track and extract the absolute minimum element across $K$ different sources in logarithmic time. This is where Heaps / Priority Queues shine.

2. Core Mental Model & Real-World Analogy

Real-Life Analogy: The Tournament of Champions

Real-Life Analogy: Imagine $K$ separate running tracks, where each track has runners starting in strictly sorted order (fastest to slowest). Instead of lining up all runners in a chaotic mess or making a single judge compare every single runner one by one, we place a scout at the starting line of each of the $K$ tracks.

The scouts report their current frontrunner to a central judge holding a Priority Queue (Min-Heap of size $K$). The judge picks the absolute fastest runner across all tracks, records them into the final sorted list, and immediately commands the track corresponding to that runner to send its next runner forward. This ensures only $K$ runners are ever evaluated at any single moment!

3. Problem 1: Merge $K$ Sorted Linked Lists (The Foundation)

Problem Statement

Given $K$ sorted linked lists where an array of size $K$ stores the root node pointers (addresses) of each linked list, our task is to merge all $K$ lists into a single, globally sorted linked list.

4. Naive Heap Approach vs. Optimized $K$-Size Heap Approach

Approach A: Dump Everything into the Heap (Naive)

• Strategy: Insert every single node from all $K$ linked lists into a single Min-Heap. Once all $N \times K$ nodes are inserted, extract them one by one.
• Time Complexity: Inserting $N \times K$ elements takes $\mathcal{O}(N K \log(NK))$, and extracting them takes another $\mathcal{O}(NK \log(NK))$. Total time: $\mathcal{O}(NK \log(NK))$.
• Space Complexity: $\mathcal{O}(NK)$ to store all nodes inside the heap.

Approach B: Restrict Heap Size to $K$ (Optimal)

• Strategy: Instead of dumping all nodes, keep the Min-Heap size capped at $K$. Initially, insert only the root node of each of the $K$ linked lists into the Min-Heap.
• Execution Flow:
  1. Extract the minimum node from the heap and append it to our result list.
  2. Check if the extracted node has a next pointer. If it does, push that next node into the Min-Heap.
  3. Repeat until the heap is empty.
• Time Complexity: $\mathcal{O}(NK \log K)$, because the heap size never exceeds $K$, making every insertion and deletion cost $\mathcal{O}(\log K)$ across all $NK$ nodes.
• Space Complexity: $\mathcal{O}(K)$ for the heap.

5. Custom Comparator & Node Pointer Mechanics in C++

When implementing Approach B in C++, we face a critical technical challenge: What do we store inside the Priority Queue?

If we store raw integer values, we lose the reference (address) to the linked list node, making it impossible to fetch the subsequent next node. Therefore, we must store Node Pointers (ListNode*).

Why a Custom Comparator is Essential

By default, a Min-Heap of ListNode* would compare the memory addresses (e.g., comparing address 100 with 700), which is completely incorrect and leads to chaotic sorting. We must write a Custom Comparator to instruct the Priority Queue to sort nodes based on their integer data values, not their memory addresses.

// Linked List Node Structure
class Node {
public:
    int data;
    Node* next;
    Node(int val) {
        data = val;
        next = NULL;
    }
};

// Custom Comparator Class
class Compare {
public:
    bool operator()(Node* a, Node* b) {
        // Standard Min-Heap logic: we want the smaller data on top.
        // NOTE: In C++ priority_queue, the comparator works inversely to sorting.
        // Returning a > b creates a Min-Heap based on node data.
        return a->data > b->data;
    }
};

Note: In C++ STL priority_queue, the comparison logic is inverted compared to standard sorting. A return value of a->data > b->data ensures that the element with the smallest data bubbles up to the top (Min-Heap behavior).

6. Complete Code Walkthrough: Merge $K$ Sorted Linked Lists

Here is the robust, production-grade C++ solution for merging $K$ sorted linked lists using a customized Min-Heap:

#include <vector>
#include <queue>

class Solution {
private:
    // Custom Comparator to maintain a Min-Heap based on Node data values
    class Compare {
    public:
        bool operator()(Node* a, Node* b) {
            return a->data > b->data;
        }
    };

public:
    Node* mergeKLists(std::vector<Node*>& lists) {
        // Define Priority Queue storing Node pointers with our Custom Comparator
        std::priority_queue<Node*, std::vector<Node*>, Compare> minHeap;

        // Step 1: Push the root node of each non-empty linked list into the heap
        for (int i = 0; i < lists.size(); i++) {
            if (lists[i] != NULL) {
                minHeap.push(lists[i]);
            }
        }

        // Step 2: Create a dummy root node to simplify result list attachment
        Node* root = new Node(0);
        Node* tail = root;

        // Step 3: Extract minimums and replenish heap from the same list
        while (!minHeap.empty()) {
            Node* temp = minHeap.top();
            minHeap.pop();

            // Attach to our output list
            tail->next = temp;
            tail = tail->next;

            // If the extracted node has a successor, push it into the heap
            if (temp->next != NULL) {
                minHeap.push(temp->next);
            }
        }

        // Return the merged list starting after the dummy node
        Node* result = root->next;
        delete root; // Clean up dummy node
        return result;
    }
};

7. Problem 2: Merge $K$ Sorted Arrays

Problem Statement

Given a 2D matrix or vector containing $K$ sorted arrays, where each array is of size $N$ (or total elements equal to $K \times N$), merge them into a single globally sorted array.

Analytical Strategy Comparison

8. Master-Level Approach: Divide & Conquer ($K$-Way Merge Sort)

Instead of using a heap, we can treat the $K$ sorted arrays like subarrays in a Merge Sort algorithm.

The Recursive Partitioning Dilemma & Solution

When dividing $K$ arrays, standard midpoint division (start + end / 2) can misalign array boundaries if array sizes or counts are odd. We manage this precisely by tracking Portions:

• At each step, we divide the active range of arrays into two halves based on the number of active portions.
• Once a portion size reduces down to 1 (meaning we are looking at a single sorted array), recursion terminates, and we merge adjacent sorted arrays linearly using two-pointer merging logic.

9. Complete Code Walkthrough: Merge $K$ Sorted Arrays (Divide & Conquer)

Here is the advanced Divide & Conquer implementation for merging $K$ sorted arrays:

class Solution {
private:
    void merge(std::vector<int>& arr, int start, int mid, int end) {
        std::vector<int> temp;
        int left = start;
        int right = mid + 1;

        while (left <= mid && right <= end) {
            if (arr[left] <= arr[right]) {
                temp.push_back(arr[left++]);
            } else {
                temp.push_back(arr[right++]);
            }
        }

        while (left <= mid) {
            temp.push_back(arr[left++]);
        }

        while (right <= end) {
            temp.push_back(arr[right++]);
        }

        for (int i = start; i <= end; i++) {
            arr[i] = temp[i - start];
        }
    }

    void mergeSortHelper(std::vector<int>& arr, int start, int end) {
        if (start >= end) return;

        int mid = start + (end - start) / 2;
        mergeSortHelper(arr, start, mid);
        mergeSortHelper(arr, mid + 1, end);
        merge(arr, start, mid, end);
    }

public:
    std::vector<int> mergeKArrays(std::vector<std::vector<int>>& arrays, int k) {
        std::vector<int> combined;
        
        // Flatten input arrays into a single container
        for (int i = 0; i < k; i++) {
            for (int j = 0; j < arrays[i].size(); j++) {
                combined.push_back(arrays[i][j]);
            }
        }

        // Apply specialized merge sort
        if (!combined.empty()) {
            mergeSortHelper(combined, 0, combined.size() - 1);
        }

        return combined;
    }
};

Summary and Key Takeaways

• Heap Size Bounding: For $K$-way merging problems, never dump all elements into a heap; cap the heap size strictly to $K$ to achieve optimal $\mathcal{O}(N K \log K)$ performance.
• Pointer Preservation: When merging Linked Lists in a heap, store Node* pointers rather than primitive integers so you can traverse to subsequent nodes dynamically.
• Custom Comparators in C++: Use custom operator structs (class Compare { public: bool operator()(Node* a, Node* b) { return a->data > b->data; } }) to override default memory-address comparisons with custom data-value logic.
• Min-Heap Inversion Rule: Remember that C++ STL priority_queue reverses logical operator expectations; returning a > b constructs a Min-Heap.
• Dummy Node Utility: Introduce a dummy root node (new Node(0)) when constructing linked lists to eliminate edge-case conditional checks for the head pointer.
• Alternative Paradigms: Beyond heaps, $K$-way merge problems can be solved with Divide & Conquer (Merge Sort partitioning), showcasing elite algorithmic versatility during technical interviews.