Merge K sorted arrays

The problem, simply

You get k arrays. Each one is sorted from small to big. Make one array that holds every number from all of them, sorted from small to big. Let N be the total count of numbers.

For [[1, 5, 9], [2, 3], [4, 8, 10]] the answer is [1, 2, 3, 4, 5, 8, 9, 10].

A story to remember it

Picture it

Examples

Way 1: the simple way

Throw every number into one big array and sort it. This ignores that the arrays were sorted.

1. Copy all N numbers into one array.
2. Sort it.

import java.util.Arrays;

class Solution {
    public int[] mergeKArrays(int[][] arrays) {
        int total = 0;
        for (int[] a : arrays) total += a.length;
        int[] out = new int[total];
        int pos = 0;
        for (int[] a : arrays) {
            for (int x : a) out[pos++] = x;
        }
        Arrays.sort(out);
        return out;
    }
}

How long it takes: sorting N numbers is O(N log N). Extra room: the output, O(N), which you need anyway.

Way 2: the clever way

Do what the lunch lady does. The heap holds at most one number per array, along with which array it came from and its position there.

1. Put the first number of every non-empty array into a min heap as {value, array, position}.
2. While the heap is not empty:
   • Poll the smallest entry and write its value to the output.
   • If its array has a next number, offer that next number.
3. Return the output.

import java.util.PriorityQueue;

class Solution {
    public int[] mergeKArrays(int[][] arrays) {
        int total = 0;
        for (int[] a : arrays) total += a.length;
        int[] out = new int[total];

        // each entry: {value, which array, position in that array}
        PriorityQueue<int[]> heap = new PriorityQueue<>((x, y) -> Integer.compare(x[0], y[0]));
        for (int i = 0; i < arrays.length; i++) {
            if (arrays[i].length > 0) heap.offer(new int[] { arrays[i][0], i, 0 });
        }

        int pos = 0;
        while (!heap.isEmpty()) {
            int[] top = heap.poll();                  // smallest front number
            out[pos++] = top[0];
            int a = top[1], next = top[2] + 1;
            if (next < arrays[a].length) {
                heap.offer(new int[] { arrays[a][next], a, next }); // its array sends the next one
            }
        }
        return out;
    }
}

How long it takes: each of the N numbers is offered once and polled once on a heap of at most k entries, so O(N log k). Extra room: the heap, O(k), plus the output.

Another way: merge in pairs

Merge array 1 with array 2, array 3 with array 4, and so on, like the rounds of a tournament. Then merge the winners in pairs again. There are log k rounds and each round touches all N numbers, so this is also O(N log k), with no heap at all. Merging one array after another into a growing result is the slow version: it is O(N k), because the early numbers get copied again in every merge.

class Solution {
    public int[] mergeKArrays(int[][] arrays) {
        if (arrays.length == 0) return new int[0];
        return mergeRange(arrays, 0, arrays.length - 1);
    }

    private int[] mergeRange(int[][] arrays, int lo, int hi) {
        if (lo == hi) return arrays[lo];
        int mid = (lo + hi) / 2;
        return mergeTwo(mergeRange(arrays, lo, mid), mergeRange(arrays, mid + 1, hi));
    }

    private int[] mergeTwo(int[] a, int[] b) {
        int[] out = new int[a.length + b.length];
        int i = 0, j = 0, p = 0;
        while (i < a.length && j < b.length) out[p++] = a[i] <= b[j] ? a[i++] : b[j++];
        while (i < a.length) out[p++] = a[i++];
        while (j < b.length) out[p++] = b[j++];
        return out;
    }
}

Watch out for

• Skip empty arrays when you fill the heap at the start, or arrays[i][0] throws.
• Compare with Integer.compare(x[0], y[0]). Writing x[0] - y[0] overflows when one number is very big and the other very negative.
• The heap must never hold two numbers from the same array. If it did, it would grow past k and lose the log k speed.
• N and k are different. Say which is which when you give the complexity.

What interviewers ask next

• Merge k sorted linked lists? The same heap, but each entry is a list node, and the next one is node.next.
• Data too big for memory (external sort)? Sort chunks that fit, write each to disk, then merge all chunks with this heap, reading each file in order. This is how databases sort huge tables.
• Smallest range that covers at least one number from each of the k lists? The same heap of front numbers, plus the current maximum.
• K-th smallest number in a matrix whose rows are sorted? Run this merge and stop after k polls, O(k log rows).
• Heap or pairwise merging? Both are O(N log k). The heap works on streams that arrive over time. Pairwise merging has no heap overhead and is easy to run in parallel.

Remember this

Grown-up words: the lunch lines are sorted arrays, the stage is a min heap (priority queue) of size k, calling the shortest kid is poll, the next kid stepping up is offer, and the whole thing is a k-way merge.