Merge K sorted Arrays

Problem Link

Merge K Sorted Arrays – GeeksforGeeks

Problem: Merge K Sorted Arrays

Problem Statement

You are given K sorted arrays, each of size K. Your task is to merge all the arrays into a single sorted array of size K*K.

Examples

Example 1
Input:

K = 3  
arr[][] = {  
    {1, 2, 3},  
    {4, 5, 6},  
    {7, 8, 9}  
}

Output: [1, 2, 3, 4, 5, 6, 7, 8, 9]

Example 2
Input:

K = 2  
arr[][] = {  
    {5, 8},  
    {1, 3}  
}

Output: [1, 3, 5, 8]

Constraints

• 1 <= K <= 100

• 1 <= arr[i][j] <= 10⁴

Intuition

Each row of the 2D array is already sorted. To merge all these sorted arrays efficiently, we can use a min-heap (priority queue). At any moment, we want to know the smallest among the remaining elements of all arrays.

By keeping only the front element of each array in the heap, we can always pick the smallest one and push the next element from that array, similar to how merge sort combines two sorted arrays.

Approach

1. Use a min-heap to store elements in the form of (value, row, col) – where:

   • value is the current number,

   • row is the index of the array it came from,

   • col is the index within that row.

2. Initialize the heap by pushing the first element of each row (i.e., arr[i][0]) into the heap.

3. While the heap is not empty:

   • Extract the smallest element from the heap and add it to the result list.

   • If there is another element in the same row (i.e., col + 1 < K), push that into the heap.

4. Continue this until all elements are merged into the result list.

Code (Java)

import java.util.*;

class Triplet {
    int value;
    int row;
    int col;

    Triplet(int value, int row, int col) {
        this.value = value;
        this.row = row;
        this.col = col;
    }
}

class Solution {
    public static ArrayList<Integer> mergeKArrays(int[][] arr, int K) {
        PriorityQueue<Triplet> pq = new PriorityQueue<>((a, b) -> a.value - b.value);
        ArrayList<Integer> result = new ArrayList<>();

        // Push the first element of each array
        for (int i = 0; i < K; i++) {
            pq.add(new Triplet(arr[i][0], i, 0));
        }

        // Extract the minimum and push the next element from the same array
        while (!pq.isEmpty()) {
            Triplet curr = pq.poll();
            result.add(curr.value);

            if (curr.col + 1 < K) {
                pq.add(new Triplet(arr[curr.row][curr.col + 1], curr.row, curr.col + 1));
            }
        }

        return result;
    }
}

Time and Space Complexity

• Time Complexity: O(K² * log K)
  There are K² elements to process, and each insertion/deletion in the heap of size K takes O(log K) time.

• Space Complexity: O(K)
  The heap stores at most K elements at any time. The output takes O(K²) space.

Dry Run

Input:

arr = {
    {1, 4, 7},
    {2, 5, 8},
    {3, 6, 9}
}, K = 3

• Initial heap: [1(from row 0), 2(from row 1), 3(from row 2)]

• Extract 1 → push 4 from row 0

• Extract 2 → push 5 from row 1

• Extract 3 → push 6 from row 2

• Continue this until all values are merged in sorted order.

Final result: [1, 2, 3, 4, 5, 6, 7, 8, 9]

Conclusion

This is a classic heap-based merging problem where a min-heap is used to efficiently combine multiple sorted arrays. It is widely used in external sorting and streaming large sorted data sources. The solution mimics the behavior of merging in K-way merge sort.