how to place a linked list node from a list of linked lists onto a heap in python

Question

I think the title is more confusing than the actual problem description. So here it is

I have a python list of linked lists, which looks somethin like lists = [[1,4,5],[1,3,4],[2,6]], and a linked list node is defined as

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

In the problem I'm trying to solve, the first step is to put a first element of each list onto a heap. Here is my code for it:

mylist1 = ListNode(1, ListNode(4, ListNode(5)))
mylist2 = ListNode(1, ListNode(3, ListNode(4)))
lists = [mylist1, mylist2]

import heapq
heap = []
for l in lists:
    heapq.heappush(heap, (l.val, l))

When I run this code I'm getting a following exception:

Traceback (most recent call last):
  File "<pyshell#29>", line 2, in <module>
    heapq.heappush(heap, (l.val, l))
TypeError: '<' not supported between instances of 'ListNode' and 'ListNode'

I think I understand what is going on there, python treats l not as a single node, but as a whole list. So my question is, how do I put a first node from each list onto a heap in python? I feel like the answer is quite simple, but I'm missing something important.

EDIT

I'm trying to solve this problem, https://leetcode.com/problems/merge-k-sorted-lists/. The approach is to use a heap, but I don't want to insert all the elements into it, what I want to do is to insert the first elements of each list, find the minimum element, then insert next element from the list which has the minimum and so on until I go through all the elements.

Answer 1

The problem with putting (l.val, l) on the heap is that when two tuples have the same value in the first member (val), that then a comparison will happen between the second member values, i.e. between lists. But this operation is not defined. You can avoid this, by putting a third, unique value in those tuples, between the two members you already have. So you could define your heap like this:

for i, l in enumerate(lists):
    heapq.heappush(heap, (l.val, i, l))

As lists can be empty, you should add a condition:

for i, l in enumerate(lists):
    if l:
        heapq.heappush(heap, (l.val, i, l))

This is also the moment to tell you that it is more efficient to first collect those tuples in a standard list, and then call heapify on it:

heap = [(l.val, i, l) for i, l in enumerate(lists) if l]
heapq.heapify(heap);

Not your question, but if you cannot make the complete logic to work, here is a spoiler for the rest of the code:

tail = dummy = ListNode() while heap: _, i, tail.next = heap[0] tail = tail.next if tail.next: heapq.heapreplace(heap, (tail.next.val, i, tail.next)) else: heapq.heappop(heap) return dummy.next