Merge K Sorted Lists

https://www.pepcoding.com/resources/online-java-foundation/hashmap-and-heap/merge-k-sorted-lists-official/ojquestion

Sample Input:

4
5
10 20 30 40 50
7
5 7 9 11 19 55 57
3
1 2 3
2
32 39

Sample Output: 1 2 3 5 7 9 10 11 19 20 30 32 39 40 50 55 57

package pep.Day65;

import java.io.*;
import java.util.*;

public class Merge_K_Sorted_Lists {

    /*
4
5
10 20 30 40 50
7
5 7 9 11 19 55 57
3
1 2 3
2
32 39

     */

    static class Pair implements Comparable<Pair> {
        int value, listIndex, indexInList;

        Pair(int val, int listIndex, int indexInList) {
            this.value = val;
            this.listIndex = listIndex;
            this.indexInList = indexInList;
        }

        @Override
        public int compareTo(Pair other) {
            return this.value - other.value;
        }
    }

    public static ArrayList<Integer> mergeKSortedLists(ArrayList<ArrayList<Integer>> lists) {
        ArrayList<Integer> ansList = new ArrayList<>();

        // write your code here
        PriorityQueue<Pair> pq = new PriorityQueue<>();

        // sab list se ek ek element priorityQueue me add kr diya
        for (int i = 0; i < lists.size(); i++) {
            pq.add(new Pair(lists.get(i).get(0), i, 0));
        }

        while (!pq.isEmpty()) {
            Pair out = pq.remove();
            ansList.add(out.value);
            out.indexInList++;

            // jis list se element remove hua hai, usme check krenge ki aur bhi element hain uss list me
            if (out.indexInList < lists.get(out.listIndex).size()) {
                pq.add(new Pair(lists.get(out.listIndex).get(out.indexInList), out.listIndex, out.indexInList));
            }
        }


        return ansList;
    }

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int k = Integer.parseInt(br.readLine());
        ArrayList<ArrayList<Integer>> lists = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            ArrayList<Integer> list = new ArrayList<>();

            int n = Integer.parseInt(br.readLine());
            String[] elements = br.readLine().split(" ");
            for (int j = 0; j < n; j++) {
                list.add(Integer.parseInt(elements[j]));
            }

            lists.add(list);
        }

        ArrayList<Integer> mlist = mergeKSortedLists(lists);
        for (int val : mlist) {
            System.out.print(val + " ");
        }
        System.out.println();
    }

}