Merge K Sorted Arrays in Java

by

This is a classic interview question. Another similar problem is “merge k sorted lists“.

This problem can be solved by using a heap. The time complexity is O(nlog(k)), where n is the total number of elements and k is the number of arrays.

It takes O(log(k)) to insert an element to the heap and it takes O(log(k)) to delete the minimum element.

class ArrayContainer implements Comparable<ArrayContainer> {
	int[] arr;
	int index;
 
	public ArrayContainer(int[] arr, int index) {
		this.arr = arr;
		this.index = index;
	}
 
	@Override
	public int compareTo(ArrayContainer o) {
		return this.arr[this.index] - o.arr[o.index];
	}
}
public class KSortedArray {
	public static int[] mergeKSortedArray(int[][] arr) {
		//PriorityQueue is heap in Java 
		PriorityQueue<ArrayContainer> queue = new PriorityQueue<ArrayContainer>();
		int total=0;
 
		//add arrays to heap
		for (int i = 0; i < arr.length; i++) {
			queue.add(new ArrayContainer(arr[i], 0));
			total = total + arr[i].length;
		}
 
		int m=0;
		int result[] = new int[total];
 
		//while heap is not empty
		while(!queue.isEmpty()){
			ArrayContainer ac = queue.poll();
			result[m++]=ac.arr[ac.index];
 
			if(ac.index < ac.arr.length-1){
				queue.add(new ArrayContainer(ac.arr, ac.index+1));
			}
		}
 
		return result;
	}
 
	public static void main(String[] args) {
		int[] arr1 = { 1, 3, 5, 7 };
		int[] arr2 = { 2, 4, 6, 8 };
		int[] arr3 = { 0, 9, 10, 11 };
 
		int[] result = mergeKSortedArray(new int[][] { arr1, arr2, arr3 });
		System.out.println(Arrays.toString(result));
	}
}

16 thoughts on “Merge K Sorted Arrays in Java”

  2. why dont we just dont do this instead:
    public static int[] mergeKSortedArray(int[][] arr) {
    PriorityQueue queue = new PriorityQueue();
    for(int i=0; i<arr.length; i++) {
    for(int j=0; j<arr[i].length; j++) {
    queue.add(arr[i][j]);
    }
    }
    int[] result = new int[queue.size()];
    int index = 0;
    while(!queue.isEmpty()) {
    result[index++] = queue.poll();
    }
    return result;
    }

    We can avoid creating the ArrayContainer class. compare to by default does the natural ordering (ascending order)

  3. Java program for Merge n number of sorted arrays.

    public class MergeSotedArrays {
    public static void main(String[] args) {
    int[] a = { 1, 4, 5 };
    int[] b = { 2, 3, 6, 7 };
    int[] c = { 1, 8, 2, 10 };
    int[] d = { 3, 11, 50, 60 };
    int[] result = merge(d,a,c,b);

    for (int idx = 0; idx < result.length; idx++) {
    System.out.print(result[idx] + " ");
    }
    }

    public static int[] merge(int[] …inputArrays) {
    int[] finalResult = inputArrays[0];
    for (int arrayno = 1; arrayno < inputArrays.length; arrayno++) {
    int tempResult[] = new int[finalResult.length + inputArrays[arrayno].length];
    int resultIndex = 0;
    for (int j = 0, k = 0; (j < finalResult.length || k < inputArrays[arrayno].length);) {
    if (j < finalResult.length && k < inputArrays[arrayno].length)
    tempResult[resultIndex++] = (finalResult[j] <= inputArrays[arrayno][k]) ? finalResult[j++] : inputArrays[arrayno][k++];
    else
    tempResult[resultIndex++] = (j < finalResult.length) ? finalResult[j++] : inputArrays[arrayno][k++];
    }
    finalResult = tempResult;
    }
    return finalResult;
    }
    }

  5. I asked my self the same question… in my opinion is O(n*log(k)) = O(n)… unless there is something I do not understand. So, without going into details (according to the PriorityQueue documentation), each poll is O(log(k)), each add is O(log(k)), these are called n times, where n is the sum of the lengths of all array. In conclusion I think is O(n*log(k)) = O(n). Am I missing something?

    I think there is mistake here, because if you look to: https://www.programcreek.com/2013/02/leetcode-merge-k-sorted-lists-java/ is O(n), which is correct.

  7. Removing a root from a PriorityQueue is log(n) because it has to move the rest of the elements to acount for the empty spot.

  11. The problem without ArrayContainer is that, as you described, after removing from the heap you must add the “next element from the exact array which we had root from.” Knowing what the next element in the array is or what array an element came from gets pretty ugly without something like an ArrayContainer class.

    You could instead enqueue an object that has a value, an index indicated what array you are talking about, and an index indicating the current index of this array. Alternatively, it could be done with maps, but it’s even worse.

    Ultimately, I think ArrayContainer is the way to do this that makes the most sense.

  12. //you can rewrite the compareTo method like this

    public int compareTo(ArrayHolder o) {
    return this.arr[this.index] – o.arr[o.index];
    }

  15. For example, I know we should put first elements of all arrays into the heap first, heapify it, then remove the minimum (which is root), and then we should put next element from the exact array which we had root from, to heap again, and again heapify and do this process till the end. But I cannot follow this procedure in here:

    while(!queue.isEmpty()){
    ArrayContainer ac = queue.poll();
    result[m++]=ac.arr[ac.index];

    if(ac.index < ac.arr.length-1){
    queue.add(new ArrayContainer(ac.arr, ac.index+1));
    }
    }

  16. Thank you for your great post. Could you please post a simpler solution without “ArrayContainer class”? It made it difficult to follow the solution for me. Thank you

Leave a Comment