Merge k sorted link lists in python

Question

I got asked this following interview for the following problem. Wondering if you have any thought and feedback for 3 solutions below:

Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity. Link to the problem: https://leetcode.com/problems/merge-k-sorted-lists/description/

# Time:  O(nlogk)
# Space: O(1)

# Merge k sorted linked lists and return it as one sorted list.
# Analyze and describe its complexity.

# Definition for singly-linked list.
class ListNode(object):
    def __init__(self, x):
        self.val = x
        self.next = None

    def __repr__(self):     
        if self:        
            return "{} -> {}".format(self.val, self.next)


# Merge two by two solution.
class Solution(object):
    def mergeKLists(self, lists):
        """
        :type lists: List[ListNode]
        :rtype: ListNode
        """
        def mergeTwoLists(l1, l2):
            curr = dummy = ListNode(0)
            while l1 and l2:
                if l1.val < l2.val:
                    curr.next = l1
                    l1 = l1.next
                else:
                    curr.next = l2
                    l2 = l2.next
                curr = curr.next
            curr.next = l1 or l2
            return dummy.next

        if not lists:
            return None
        left, right = 0, len(lists) - 1;
        while right > 0:
            if left >= right:
                left = 0
            else:
                lists[left] = mergeTwoLists(lists[left], lists[right])
                left += 1
                right -= 1
        return lists[0]


# Time:  O(nlogk)
# Space: O(logk)
# Divide and Conquer solution.
class Solution2:
    # @param a list of ListNode
    # @return a ListNode
    def mergeKLists(self, lists):
        def mergeTwoLists(l1, l2):
            curr = dummy = ListNode(0)
            while l1 and l2:
                if l1.val < l2.val:
                    curr.next = l1
                    l1 = l1.next
                else:
                    curr.next = l2
                    l2 = l2.next
                curr = curr.next
            curr.next = l1 or l2
            return dummy.next

        def mergeKListsHelper(lists, begin, end):
            if begin > end:
                return None
            if begin == end:
                return lists[begin]
            return mergeTwoLists(mergeKListsHelper(lists, begin, (begin + end) / 2), \
                                 mergeKListsHelper(lists, (begin + end) / 2 + 1, end))

        return mergeKListsHelper(lists, 0, len(lists) - 1)


# Time:  O(nlogk)
# Space: O(k)
# Heap solution.
import heapq
class Solution3:
    # @param a list of ListNode
    # @return a ListNode
    def mergeKLists(self, lists):
        dummy = ListNode(0)
        current = dummy

        heap = []
        for sorted_list in lists:
            if sorted_list:
                heapq.heappush(heap, (sorted_list.val, sorted_list))

        while heap:
            smallest = heapq.heappop(heap)[1]
            current.next = smallest
            current = current.next
            if smallest.next:
                heapq.heappush(heap, (smallest.next.val, smallest.next))

        return dummy.next


if __name__ == "__main__":
    list1 = ListNode(1)
    list1.next = ListNode(3)
    list2 = ListNode(2)
    list2.next = ListNode(4)

    print Solution().mergeKLists([list1, list2])

Answer 1

Bug in Solution3 (using heapq)

The items being put on the heap need to be fully sortable. Tuples are compared element by element. If the first element of two tuples compare equal, then the next element of the two tuples are compared. In your solution, the second element of the tuple is a ListNode. But no methods for comparing ListNodes have been defined. So if the first element of the tuples are equal a TypeError will be raised when the ListNodes are compared. The solution is to use 3-tuples: (value, seq_no, ListNode), in which the seq_nos are unique.

class Solution3:
    # @param a list of ListNode
    # @return a ListNode
    def mergeKLists(self, lists):
        dummy = ListNode(0)
        current = dummy

        heap = []
        for seq_no, sorted_list in enumerate(lists):
            if sorted_list:
                heapq.heappush(heap, (sorted_list.val, seq_no, sorted_list))

        while heap:
            _, seq_no, smallest = heapq.heappop(heap)
            current.next = smallest
            current = current.next
            if smallest.next:
                heapq.heappush(heap, (smallest.next.val, seq_no, smallest.next))

        return dummy.next

See discussion in heapq documentation: https://docs.python.org/3.7/library/heapq.html#priority-queue-implementation-notes