Untitled
import java.util.*; public class MergeKSortedArrays { static class HeapNode { List<Integer> list; int index; public HeapNode(List<Integer> list, int index) { this.list = list; this.index = index; } } public static List<Integer> mergeKSortedArrays(List<List<Integer>> arrays) { // Min-Heap with custom comparator using a lambda PriorityQueue<HeapNode> minHeap = new PriorityQueue<>( (a, b) -> Integer.compare(a.list.get(a.index), b.list.get(b.index)) ); List<Integer> result = new ArrayList<>(); // Step 1: Initialize the heap with the first element of each array for (List<Integer> list : arrays) { if (!list.isEmpty()) { minHeap.add(new HeapNode(list, 0)); } } // Step 2: Process the heap while (!minHeap.isEmpty()) { // Extract the smallest element HeapNode smallest = minHeap.poll(); result.add(smallest.list.get(smallest.index)); // Add the next element from the same list to the heap if (smallest.index + 1 < smallest.list.size()) { minHeap.add(new HeapNode(smallest.list, smallest.index + 1)); } } return result; } public static void main(String[] args) { Scanner sc = new Scanner(System.in); int T = sc.nextInt(); // Number of test cases while (T-- > 0) { int K = sc.nextInt(); // Number of arrays List<List<Integer>> arrays = new ArrayList<>(); for (int i = 0; i < K; i++) { int N = sc.nextInt(); // Size of the current array List<Integer> array = new ArrayList<>(); for (int j = 0; j < N; j++) { array.add(sc.nextInt()); } arrays.add(array); } // Merge and print the result List<Integer> merged = mergeKSortedArrays(arrays); for (int num : merged) { System.out.print(num + " "); } System.out.println(); } sc.close(); } }
Leave a Comment