Merging k sorted lists using heapq module in python3

Question

Problem:- merge k sorted lists.

I want to solve this problem using min heap which can be implemented with heapq module in python. Below is the sample code of the function...

heapq.heappush(listwithoutNone,(node.val, node))
        

But the problem is that python interpreter is raising an error:

TypeError: '<' not supported between instances of 'ListNode' and 'ListNode' heapq.heappush(listwithoutNone,(node.val, node))

So, I wanna use the node.val as a minheap node element, since I am passing a tuple, so what change should I do in the code so that minheap will heapify the heap with node.val.

Thanks in Advance.

Answer 1

When tuples are compared, their first elements are compared, and then any ties are broken using their second elements, their elements, and so on. For example, (2, "a") < (2, "b") will evaluate to True.

Here, you are inserting (node.val, node) tuples into your heap, which attempts to compare them. If there is a tie on the node value, it moves on to the second element of the tuples - the nodes themselves. These are just ListNode instances. Python does not know how to compare two ListNodes, hence the error.

To enable comparisons between ListNodes, you need to implement the rich comparison methods. A quick way to do it is to simply implement ListNode.__lt__ and then use the functools.total_ordering decorator:

import heapq
from functools import total_ordering


@total_ordering
class ListNode:
    def __init__(self, value: float, label: str) -> None:
        self.value = value
        self.label = label

    def __lt__(self, other: 'ListNode'):
        return self.value <= other.value

    def __str__(self):
        return f"ListNode(label={self.label}, value={self.value})"



nodes = []
a =  ListNode(5, "A")
b = ListNode(3, "B")
c = ListNode(5, "C")
heapq.heappush(nodes, a)
heapq.heappush(nodes, b)
heapq.heappush(nodes, c)

while nodes:
    x = heapq.heappop(nodes)
    print(x)

Here we say that comparing two ListNode objects is the same as just comparing their values. With comparisons defined, you don't even need to insert tuples at all. You can just insert ListNode objects directly, and rely on the comparison methods.

Answer 2

I think you are talking about this: Leetcode Merge k Sorted Lists

I am sharing a working solution for you:

head = curr = ListNode(0)    # Creating a dummy node
    lst = []
    for l in lists:
        if l:
 # Here we need to first push val so that heap can know which is minimum and which is maximum. 
 # If we insert directly node then heap can't work proper way (in this case heapq.heappop doesn't return minimum node).    
            lst.append((l.val,l))    

    
    heapq.heapify(lst)
    while lst:
        val,node = heapq.heappop(lst)
        curr.next = node
        curr = curr.next
        node = node.next
        if node:
            heapq.heappush(lst,(node.val,node))
    return head.next 

Answer 3

To improve upon Saurabh's answer, you can generate the final sorted list while iterating over the individual lists.

The idea is to first populate the heap with the first element of each sorted array tagged with the array the element came from, and whenever popping from the heap, add the next element of the array associated with the element just popped from the heap.

def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
    heap = []
    for i, l in enumerate(lists):
        if l is None:
            continue
        heapq.heappush(heap, (l.val,i))
        lists[i] = lists[i].next

    out = ListNode() # dummy
    cur = out
    while heap:
        val, i = heapq.heappop(heap)
        cur.next = ListNode(val)
        cur = cur.next
        if lists[i] is not None:
            heapq.heappush(heap, (lists[i].val, i))
            lists[i] = lists[i].next
    return out.next

This can also be modified so that the heap keeps track of the actual nodes so that no new nodes (besides the dummy head) are created:

def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
    heap = []
    for i, l in enumerate(lists):
        if l is None:
            continue
        heapq.heappush(heap, (l.val, i, l))
        lists[i] = lists[i].next

    out = ListNode() # dummy
    cur = out
    while heap:
        val, i, node = heapq.heappop(heap)
        cur.next = node
        cur = cur.next
        if lists[i] is not None:
            heapq.heappush(heap, (lists[i].val, i, lists[i]))
            lists[i] = lists[i].next
    return out.next

The heap will only ever contain a maximum k elements where k is the number of input lists. Thus, we address @Abhishek Jaiswal's concern that the space complexity should be O(k).

Edit: add space complexity improvement analysis

Answer 4

You will need to override the lt function that ListNode offers. Here is the code

def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
    def my_lt(self, other):
        return self.val<other.val
    setattr(ListNode, "__lt__", my_lt)
    heap = []
    for l in lists:
        if l:
            heap.append((l.val, l))
    heapq.heapify(heap)
    head = ListNode(0)
    curr = head

    while heap:
        val, node = heapq.heappop(heap)
        curr.next = node
        curr = curr.next
        if node.next:
            heapq.heappush(heap, (node.next.val, node.next))
    
    return head.next