fork(1) download
  1. # Merge K Sorted Lists (LC version 3.01).
  2.  
  3. from heapq import heappush, heappop
  4.  
  5. # List.
  6.  
  7. class ListNode:
  8. def __init__(self, val=None, next=None):
  9. self.val = val
  10. self.next = next
  11.  
  12. def makeList(seq):
  13. t = None
  14. for x in reversed(seq):
  15. t = ListNode(x, t)
  16. return t
  17.  
  18. def printList(l):
  19. a = []
  20. while l:
  21. a.append(l.val)
  22. l = l.next
  23. print('List(', a, ')', sep='')
  24.  
  25. # Merge.
  26.  
  27. def mergeKLists(lists):
  28. heap = []
  29. for k, n in enumerate(lists):
  30. if n:
  31. heappush(heap, (n.val, k, n))
  32. temp = ListNode()
  33. if heap:
  34. t = temp
  35. while True:
  36. _, k, n = heappop(heap)
  37. t.next = n
  38. if not heap:
  39. break
  40. t = t.next
  41. n = n.next
  42. if n:
  43. heappush(heap, (n.val, k, n))
  44. return temp.next
  45.  
  46. # Show.
  47.  
  48. l0 = makeList([])
  49. l1 = makeList([1, 3, 5, 7])
  50. l2 = makeList([2, 4, 6])
  51. l3 = makeList([0, 4, 8])
  52.  
  53. printList(l0)
  54. printList(l1)
  55. printList(l2)
  56. printList(l3)
  57.  
  58. printList(mergeKLists([]))
  59. printList(mergeKLists([l0]))
  60. printList(mergeKLists([l1]))
  61. printList(mergeKLists([l1, l2, l3]))
Success #stdin #stdout 0.09s 14116KB
stdin
Standard input is empty
stdout
List([])
List([1, 3, 5, 7])
List([2, 4, 6])
List([0, 4, 8])
List([])
List([])
List([1, 3, 5, 7])
List([0, 1, 2, 3, 4, 4, 5, 6, 7, 8])