leetcode/MergeSortedLists_ListVer.py at master · tnstar/leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
'''
@author: TNSTAR
'''
import heapq

class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None
class Solution:
    # @param a list of ListNode
    # @return a ListNode
    def mergeKLists(self, lists):
        head = None
        phead = None
        next_pop = -1
        avail_list = range(len(lists)) #lists to be extracted
        pri_heap = []
        for i in avail_list:
            if lists[i]:
                heapq.heappush(pri_heap, (lists[i].pop(0),i))
            else:
                avail_list.remove(i)
        while (avail_list or pri_heap):
            if head == None:
                temp = heapq.heappop(pri_heap)
                #print temp[0]
                head = ListNode(temp[0])
                phead = head
                next_pop = temp[1]
                #print str(next_pop)+" next"
            else:
                temp = heapq.heappop(pri_heap)
                #print temp[0]
                head.next = ListNode(temp[0])
                head = head.next
                next_pop = temp[1]
                #print str(next_pop)+" next"
            if lists[next_pop]:
                heapq.heappush(pri_heap, (lists[next_pop].pop(0),next_pop))
                #print str(next_pop)+" pop"
                if not lists[next_pop]:
                    #print "removing"+str(next_pop)
                    avail_list.remove(next_pop)
            else:
                pass
        return phead

lists = [[1,5,8],[2,4,8,9],[3,6,9]]
sol = Solution()
head = sol.mergeKLists(lists)
while head.next!=None:
    print head.val
    head=head.next