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])

• Gareth Rees: Python's standard library has heapq.merge and the implemention is worth a look.
• NinjaG: thanks for letting me know. Wondering if that is allowed in technical interview?
• Gareth Rees: It can't hurt to mention it. "In real code I would use the built-in heapq.merge, but as an exercise I would implement it like this ..."
• Bailey Parker: Why did you put if self in __repr__? That will always be true.

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