Efficiently Merging K Sorted Arrays with Heaps

1. The Problem: Merging K Sorted Arrays

Imagine you are given K separate arrays, each already sorted in ascending order. Your task is to combine them into a single, large array that is also sorted.

A straightforward idea might be to concatenate all K arrays into one massive list and then sort it. If N is the total number of elements across all arrays, this approach would have a time complexity of O(N log N). However, this method completely ignores a crucial piece of information: the input arrays are already sorted. This suggests we can devise a more efficient strategy.

2. From Two Arrays to K: Building the Intuition

Let's simplify. How would you merge two sorted arrays? You'd likely use a two-pointer approach: compare the elements at the current pointers of each array, take the smaller one, and advance that array's pointer. This is the core logic of the merge step in Merge Sort.

Now, how can we generalize this to K arrays? The core idea remains the same: at each step, we need to find the single smallest element among all the current candidate elements (one from each array) and add it to our result.

The main challenge is this: how do we efficiently find the minimum among K elements, over and over again? A simple loop to scan the K candidates would take O(K) time for each element we select. Since we have to do this for all N elements, the total time complexity would be O(N * K), which is often slower than the naive "concatenate and sort" method.

This is where we need a better data structure to track the K candidates.

As the video suggests, the perfect tool for repeatedly and efficiently finding the minimum element from a collection is a min-heap.

3. The Min-Heap Solution

A min-heap can store the K candidate elements (one from each array) and allow us to extract the overall minimum in O(log K) time. This is a huge improvement over the O(K) linear scan.

Here is the algorithm:

1. Initialize: Create a min-heap and an empty list for the final result.
2. Seed the heap: For each of the K arrays, if it's not empty, push its first element into the heap. To keep track of where each element came from, we'll store a tuple: (value, array_index, element_index). The heap will automatically sort these tuples based on the value.
3. Process: While the heap is not empty:
   a. Pop: Extract the tuple with the smallest value from the heap. Let's call it (val, arr_idx, elem_idx).
   b. Append: Add val to our result list.
   c. Replace: Look back at the array from which the element came (arrays[arr_idx]). If there's a next element in that array, push it onto the heap as a new tuple: (next_value, arr_idx, elem_idx + 1).

This process continues until the heap is empty, which means we have processed all elements from all K arrays.

Let's watch a detailed walkthrough of this algorithm. The following video uses "linked lists" instead of arrays, but the core logic is identical.

4. Implementation and Complexity

Now, let's look at a clean Python implementation for merging K sorted arrays.

import heapq

def merge_k_sorted_arrays(arrays):
    """
    Merges K sorted arrays into a single sorted array.

    Args:
        arrays: A list of sorted lists (arrays).

    Returns:
        A single sorted list containing all elements.
    """
    min_heap = []
    
    # 1. Initialize the heap with the first element of each non-empty array.
    for i, arr in enumerate(arrays):
        if arr:
            # Push a tuple: (value, array_index, element_index_in_array)
            heapq.heappush(min_heap, (arr[0], i, 0))
            
    result = []
    
    # 2. Process elements until the heap is empty.
    while min_heap:
        # Pop the element with the smallest value.
        val, arr_idx, elem_idx = heapq.heappop(min_heap)
        
        # Add it to the result list.
        result.append(val)
        
        # If there's a next element in the array we just took from,
        # push it to the heap.
        if elem_idx + 1 < len(arrays[arr_idx]):
            next_elem_idx = elem_idx + 1
            next_val = arrays[arr_idx][next_elem_idx]
            heapq.heappush(min_heap, (next_val, arr_idx, next_elem_idx))
            
    return result

# --- Example Usage ---
k_arrays = [
    [3, 6],
    [4, 7, 11, 12],
    [1, 4, 5]
]
merged_array = merge_k_sorted_arrays(k_arrays)
print(f"Input arrays: {k_arrays}")
print(f"Merged array: {merged_array}")
# Expected Output: [1, 3, 4, 4, 5, 6, 7, 11, 12]

Complexity Analysis

Let's formalize the efficiency of this algorithm.

• K is the number of sorted arrays.

• N is the total number of elements across all arrays.

• Time Complexity: O(N log K)

  • We process each of the N elements exactly once.
  • For each element, we perform one heappop and at most one heappush.
  • The size of the heap is at most K, so each heap operation takes O(log K) time.
  • Therefore, the total time is N * O(log K) = O(N log K). This is significantly better than O(N log N) when K is much smaller than N.

• Space Complexity: O(K)

  • The primary auxiliary space used is for the min-heap, which stores at most one element from each of the K arrays.
  • The result array takes O(N) space, but this is typically considered output space, not auxiliary space.

The video you watched previously also provides a good summary of this analysis.

Conclusion

You have successfully learned how to solve the "Merge K Sorted Arrays" problem using a min-heap. This pattern is a cornerstone of many efficient algorithms.

Key Takeaways:

• The problem involves merging K pre-sorted arrays into a single sorted array.
• A naive approach of concatenating and sorting is O(N log N) and fails to use the sorted property of the inputs.
• The optimal solution uses a min-heap of size K to efficiently track the smallest available element from each array.
• This heap-based approach achieves a time complexity of O(N log K) and an auxiliary space complexity of O(K).
• Storing tuples like (value, array_index, element_index) in the heap is crucial for tracking which element to add next.

Up Next:
In our final lesson of this module, we will formalize a technique we've touched upon before: how to implement a max-heap using Python's heapq module by negating values. This is a simple but essential trick for your competitive programming toolkit.