from
heapq
import
heappop, heapify
class
ListNode:
def
__init__(
self
, x,
next
=
None
):
self
.val
=
x
self
.
next
=
next
class
Solution:
def
mergeKLists(
self
, lists):
h
=
[]
for
node
in
lists:
while
node:
h.append(node.val)
node
=
node.
next
if
not
h:
return
None
heapify(h)
root
=
ListNode(heappop(h))
curnode
=
root
while
h:
nextnode
=
ListNode(heappop(h))
curnode.
next
=
nextnode
curnode
=
nextnode
return
root
def
to_list(head):
res
=
[]
curnode
=
head
while
curnode:
res.append(curnode.val)
curnode
=
curnode.
next
return
res
def
test_merge():
l1
=
ListNode(
1
, ListNode(
4
, ListNode(
5
)))
l2
=
ListNode(
1
, ListNode(
3
, ListNode(
4
)))
s
=
Solution()
lists
=
[l1, l2]
head
=
s.mergeKLists(lists)
print
(to_list(head))
if
__name__
=
=
'__main__'
:
test_merge()