72. Merge k Sorted Arrays

Topic :

priority queue

Difficulty :

medium

Problem Link :

problem statement

Given K sorted arrays arranged in the form of a matrix of size K*K. The task is to merge them into one sorted array.

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
Explanation:Above test case has 3 sorted
arrays of size 3, 3, 3
arr[][] = [[1, 2, 3],[4, 5, 6], 
[7, 8, 9]]
The merged list will be 
[1, 2, 3, 4, 5, 6, 7, 8, 9].

Example 2:

Input:
K = 4
arr[][]={{1,2,3,4},{2,2,3,4},
         {5,5,6,6},{7,8,9,9}}
Output:
1 2 2 2 3 3 4 4 5 5 6 6 7 8 9 9 
Explanation: Above test case has 4 sorted
arrays of size 4, 4, 4, 4
arr[][] = [[1, 2, 2, 2], [3, 3, 4, 4],
[5, 5, 6, 6], [7, 8, 9, 9 ]]
The merged list will be 
[1, 2, 2, 2, 3, 3, 4, 4, 5, 5, 
6, 6, 7, 8, 9, 9].

You do not need to read input or print anything. Your task is to complete mergeKArrays() function which takes 2 arguments, an arr[K][K] 2D Matrix containing K sorted arrays and an integer K denoting the number of sorted arrays, as input and returns the merged sorted array ( as a pointer to the merged sorted arrays in cpp, as an ArrayList in java, and list in python)

Constraints:
1 <= K <= 100

solution

import java.io.*;
import java.util.*; 
class Merge_K_sorted_arrays
{
    public static void main(String args[])
    {
        int[][]arr= {{1,2,3},{4,5,6},{7,8,9}};
        int K=3;
        System.out.println(mergeKArrays(arr,K));
    }
    
    static class Data
    {
      int array;
      int index;
      int val;
      
      Data(int array,int index,int val)
      {
          this.array=array;
          this.index=index;
          this.val=val;
      }
      
    }
    
    public static ArrayList<Integer> mergeKArrays(int[][] arr,int K) 
    {
      ArrayList<Integer> ans=new ArrayList<>();
      PriorityQueue<Data> pq=new PriorityQueue<>((a,b)->(a.val-b.val));
      
      for(int i=0;i<K;i++)
      {
          pq.add(new Data(i,0,arr[i][0]));
      }
      
      while(!pq.isEmpty())
      {
          Data curr=pq.poll();
          ans.add(curr.val);
          
          if(curr.index+1<arr[curr.array].length)
            pq.offer(new Data(curr.array,curr.index+1,arr[curr.array][curr.index+1]));
      }
      
      return ans;
    }
}