Merging k sorted lists the way merge sort does

$ cat the-problem.mdThe problem

We are given an array of $k$ linked listsLinked list A list in which each element, called a node, holds a value and a link to the next node., each sorted in ascending order. We have to merge them all into one sorted linked list and return it. Throughout, $k$ is the number of lists and $N$ is the total number of nodes across all of them.

Example 1

• Input: lists = [[2, 6], [0, 7, 9], [3, 8]]
• Output: [0, 2, 3, 6, 7, 8, 9]
• Why: these are the seven numbers of the three lists, in ascending order. Here $k = 3$ and $N = 7$.

Example 2

• Input: lists = [[1, 4, 5], [1, 3, 4]]
• Output: [1, 1, 3, 4, 4, 5]
• Why: a number that appears in more than one list appears that many times in the result.

Example 3

• Input: lists = [[], [5]]
• Output: [5]
• Why: a list can be empty. If every list is empty, or there are no lists at all, the result is an empty list.

$ cat introduction.mdIntroduction

The common way to solve this problem in $O(N \log k)$O(N log k) The work grows like N times log k. Here log k is about the number of times k can be halved before it reaches 1. For 8 lists that is 3. time is to use a heapHeap A structure that always gives back its smallest item. With k items in it, adding one or removing the smallest takes about log k steps. Here it would hold the next node of each list.. However, there is another solution. It is based on the observation that at one point in its execution, merge sortMerge sort A way of sorting a list: split it into halves, sort each half, then merge the two sorted halves into one sorted list. can be said to be merging $k > 0$ sorted lists. Therefore, we should be able to adapt merge sort to this problem. If we adapt it successfully, we should get the same time complexity of $O(N \log k)$, without a heap.

The picture below shows merge sort working on eight numbers, one level per row. It starts with every number as a sorted list of its own, and each level merges the lists of the level above in pairs. Choose a number of lists, $k$, to see the level where merge sort is holding exactly $k$ sorted lists. Everything below that level is merge sort merging $k$ sorted lists.

$ cat approach.mdApproach

Merge sort does more than merge two lists at a time. It also merges the lists in levels, and the lists in the earlier levels are the smaller ones. So, to reach the same asymptotic complexityAsymptotic complexity How the time or memory a solution needs grows as its input grows, leaving out constant factors. It is written in O notation. as merge sort, we need to:

• Merge the lists in pairs, the way bottom-up merge sortBottom-up merge sort Merge sort without the splitting. It starts from lists of one element and merges neighbours in pairs, pass after pass, until one list is left. does.

To do this, we use a queueQueue A line of items. New items join at the back and items leave from the front, so the first one in is the first one out.. We start by putting all of the lists on the queue. Then we take the first two lists, l1 and l2, off the front of the queue and merge them to produce merged. To make sure merged is processed after all the other lists already on the queue, we add merged to the end of the queue. We repeat this until only one list is left on the queue. That list is the answer.

$ cat try-it-yourself.mdTry it yourself

Press Step to merge the two lists at the front of the queue. Each list shows its level, which is the number of merges it took to build. Then switch the merged list to the front of the queue instead of the back and step through again. The lists reach deeper levels, and more nodes have to be merged.

$ cat why-the-back-of-the-queue-matters.mdWhy the back of the queue matters

The time complexity rests on one claim: no list ever becomes more than $\lceil \log_2 k \rceil$ levels deep. The queue is what guarantees it.

A merged list goes to the back of the queue, behind every list that is still waiting. Because the queue is first in, first out, the merged list cannot be merged again until all the lists ahead of it have had their turn. This keeps the levels rising together. The $k$ lists at level 0 are merged in pairs into about $k/2$ lists at level 1, and no level 1 list is merged again until every level 0 list has had its turn. Those become about $k/4$ lists at level 2, and so on. The number of lists halves with each level, so after $\lceil \log_2 k \rceil$ levels only one list is left.

When a level has an odd number of lists, one list is left over, and it is merged with a list one level deeper. It ends up exactly as deep as it would have been if it had first been merged with an empty list. So $k$ lists never go deeper than the next power of two would, which is $\lceil \log_2 k \rceil$ levels.

You can see both halves of this in the demonstration above. With the merged list going to the back, the six lists finish at level 3, which is $\lceil \log_2 6 \rceil$. With it going to the front, the newly merged list is merged again straight away, ahead of lists that have not yet had a turn. The same six lists finish at level 5, and the numbers in the first two lists are handled five times instead of three.

$ cat a-note-on-python.mdA note on Python

We could use an ordinary Python list as the queue and remove elements from its front. However, removing the first element of a Python list is not a constant-time operation: every element after it has to move up one place, so with $k$ lists on the queue it takes $O(k)$ time. To get the time complexity we want, removing from the front of the queue has to take $O(1)$ time. We can achieve this with a dequeDeque Short for double-ended queue: a queue that can add and remove items quickly at both ends. Python's is collections.deque. from Python’s collections library.

$ cat complexity.mdComplexity

Let $N$ be the total number of nodes across all lists, and let $k$ be the number of lists.

• Time complexity: $O(N \log k)$, the same as the heap solution. A merge handles each of its nodes in constant time and produces a list one level deeper. As shown above, sending merged lists to the back of the queue means no list ever gets more than $\lceil \log_2 k \rceil$ levels deep, so each node is handled at most that many times. Strictly, building the queue and making the $k - 1$ merges adds $O(k)$ on top, for $O(k + N \log k)$. The extra term only matters when most of the lists are empty, because otherwise $k \le N$.
• Space complexity: $O(k)$, for the queue. Merging relinks the nodes we were given and creates no new ones.

$ cat code.mdCode

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
from collections import deque


class Solution:
    def mergeKLists(self, lists: list[ListNode | None]) -> ListNode | None:
        if len(lists) == 1:
            return lists[0]

        queue = self.init_queue(lists)
        merged = None
        while len(queue) > 1:
            l1 = queue.popleft()
            l2 = queue.popleft()

            merged = self.merge_two_lists(l1, l2)
            queue.append(merged)

        return merged

    def merge_two_lists(self, l1, l2):
        if l1 is None:
            return l2
        if l2 is None:
            return l1

        # At this point, we know neither list is empty
        l1, l2, head, temp = self.init_head_and_temp(l1, l2)

        while l1 and l2:
            if l1.val <= l2.val:
                temp.next = l1
                l1 = l1.next
            else:
                temp.next = l2
                l2 = l2.next
            temp = temp.next

        remainder = l1 if l1 is not None else l2
        temp.next = remainder
        return head

    def init_head_and_temp(self, l1, l2):
        if l1.val <= l2.val:
            head = l1
            l1 = l1.next
        else:
            head = l2
            l2 = l2.next
        temp = head
        return (l1, l2, head, temp)

    def init_queue(self, lists):
        queue = deque()
        for l in lists:
            queue.append(l)
        return queue