实现优先级队列(小根堆)

Posted 2020-11-25 npc-assange

头条算法一面，记录一下

实现优先级队列(小根堆)

 1 class prioQueue:
 2     def __init__(self, elist=[]):
 3         self._elems = list(elist)
 4         if elist:
 5             self.buildheap()
 6     def siftdown(self, e,begin, end):
 7         elems, i, j, = self._elems, begin , begin*2+1
 8         while j < end:
 9             if j+1<end and elems[j+1] < elems[j]:
10                 j = j + 1         
11             if e < elems[j]:
12                 break
13             elems[i] = elems[j]
14             i, j = j , 2*j+1
15         elems[i] = e
16     def buildheap(self):
17         end = len(self._elems)
18         for i in range(end//2,-1,-1):
19             self.siftdown(self._elems[i],i,end)
20         return self._elems
21     def peek(self):                       
22         return self._elems[0]
23     def dequeue(self):
24         elems = self._elems
25         e0 = elems[0]
26         e = elems.pop()
27         if len(elems) > 0:
28             self.siftdown(e, 0 , len(elems))
29         return e0
30     def siftup(self, e, last):
31         elems , i , j = self._elems, last, (last-1)//2
32         while i > 0 and e < elems[j]:
33             elems[i] = elems[j]
34             i, j = j ,(j-1)//2
35         elems[i] = e
36     def enqueue(self, e):
37         self._elems.append(None)
38         self.siftup(e, len(self._elems)-1)
39         return self._elems