Merge K Sorted Arrays Editorial

Practice Problem Link: Merge K Sorted Arrays | Practice Problem

Please make sure to try solving the problem yourself before looking at the editorial.

Problem Statement

You are given k sorted arrays in the form of 2D integer matrix arr of size k*n. Merge them into a single sorted array.

Naive Approach

A simple way to solve this problem is to store all the elements in a single array and sort them.

Analysis

• Time Complexity: O(n * k * log(n * k))
• Auxiliary Space Complexity: O(1)

Implementation

C++

vector<int> mergeKArrays(vector<vector<int>> &arr) {
   vector<int> sortedArray;
   for (int i = 0; i < arr.size(); i++) {
       for (int j = 0; j < arr[0].size(); j++) {
           sortedArray.push_back(arr[i][j]);
       }
   }
   sort (sortedArray.begin(), sortedArray.end());
   return sortedArray;
}

Java

class Solution {
   int[] mergeKArrays(int[][] arr) {
       int n = arr.length * arr[0].length;
         int[] sortedArray = new int[n];
       int indx = 0;
       for (int i = 0; i < arr.length; i++) {
           for (int j = 0; j < arr[0].length; j++) {
               sortedArray[indx++] = arr[i][j];
           }
       }
       Arrays.sort(sortedArray);
       return sortedArray;
   }
}

Optimal Approach

This problem can be solved efficiently using Min Heap. 

• Since all the arrays are sorted, the first element of all K arrays can be stored in the Min heap because the first element of the sorted array must be one of those elements. 
• Remove the first element from the Min heap and add it to the answer array. Add the next element of the same row from which the element was removed, in the Min heap. Now one of these K elements must be the second element in the sorted array.
• Repeat the same process until all the elements are added to the answer array in a sorted manner.

Analysis

• Time Complexity: O(n * k * log(k))
• Auxiliary Space Complexity: O(k)

Implementation

C++

vector<int> mergeKArrays(vector<vector<int>> &arr) {
   int n = arr.size() * arr[0].size();
   priority_queue <pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>, greater<pair<int, pair<int, int>>>> minHeap;
   for (int i = 0; i < arr.size(); i++) {
       minHeap.push( {arr[i][0], {i, 0} } );
   }
   vector<int> sortedArray;
   while(!minHeap.empty()) {
       sortedArray.push_back(minHeap.top().first);
       pair<int, int> index = {minHeap.top().second.first, minHeap.top().second.second + 1};
       minHeap.pop();
       if (index.second != arr[0].size()) {
           minHeap.push({arr[index.first][index.second], {index.first, index.second}});
       }
   }
   return sortedArray;
}

Java

class HeapNode implements Comparable<HeapNode> {
   int value, row, column, size;
   public HeapNode(int value, int row, int column, int size) {
       this.value = value;
       this.row = row;
       this.column = column;
       this.size = size;
   }
   public boolean hasNext () {
       return this.column < this.size - 1;
   }
   public int compareTo (HeapNode node) {
       return value - node.value;
   }
}
class Solution {
   int[] mergeKArrays(int[][] arr) {
       int n = arr.length * arr[0].length;
       PriorityQueue<HeapNode> minHeap = new PriorityQueue<>();
       for (int i = 0; i < arr.length; i++) {
           minHeap.add(new HeapNode(arr[i][0], i, 0, arr[0].length));
       }
       int[] sortedArray = new int[n];
       int i = 0;
       while(!minHeap.isEmpty()) {
           HeapNode node = minHeap.poll();
           sortedArray[i++] = node.value;
           if(node.hasNext()) {
               minHeap.add(new HeapNode(arr[node.row][node.column + 1], node.row, node.column + 1, arr[0].length));
           }
       }
       return sortedArray;
   }
}