Merge k Sorted Lists

The problem description can be found in LeetCode.

Approach

A naive approach would consist in moving all the elements that belong to the k linked lists to an array. Then, we can simply sort it. This would require $O(nlog(n))$ time for sorting the list and $O(n)$ space for storing all the elements (assuming that $n={\sum }_{i=0}^{len(lists)-1}len(lists[i])$).

We could improve this approach by just keeping track of the minimum element of each list:

l1 = [1, 2, 3]
l2 = [5, 6, 7]
l3 = [10, 11]

At first, we could keep min_list = [1, 5, 10], which are the minimum values for l1, l2 and l3, respectively. The minimum of this list would be the first value of the result list.

Then, since we pick the minimum from the first list, we should rearrange min_list based on the next value of l1, which is 2. This results in min_list = [2, 5, 10]. Again, we can peek the first element and use it as the second element of the result list.

We can repeat this process until we have elements in our list.

Of course, the perfect data structure we can use to keep track the minimum is the Heap.

Solution

from heapq import heappop, heappush
from typing import List, Optional


class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        heap = []

        # use the `id` to avoid to define a `heap container`,
        # which needs `__lt__` to be implemented
        for head in lists:
            if head:
                heappush(heap, (head.val, id(head), head))

        head = ListNode(-1)
        p = head
        while heap:
            _, _, node = heappop(heap)
            if node.next:
                heappush(heap, (node.next.val, id(node.next), node.next))

            p.next = ListNode(node.val)
            p = p.next

        return head.next

Complexity

• Time: $O(n)$
• Space: $O(n)$

Where $n$ is the size of the nums list.