forked from nd7141/influence-maximization
-
Notifications
You must be signed in to change notification settings - Fork 0
/
priorityQueue.py
46 lines (38 loc) · 1.6 KB
/
priorityQueue.py
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
__author__ = 'http://docs.python.org/2/library/heapq.html#priority-queue-implementation-notes'
import itertools
from heapq import *
class PriorityQueue(object):
def __init__(self):
self.pq = [] # list of entries arranged in a heap
self.entry_finder = {} # mapping of tasks to entries
self.REMOVED = '<removed-task>' # placeholder for a removed task
self.counter = itertools.count() # unique sequence count
def add_task(self, task, priority=0):
'Add a new task or update the priority of an existing task'
if task in self.entry_finder:
self.remove_task(task)
count = next(self.counter)
entry = [priority, count, task]
self.entry_finder[task] = entry
heappush(self.pq, entry)
def remove_task(self, task):
'Mark an existing task as REMOVED. Raise KeyError if not found.'
entry = self.entry_finder.pop(task)
entry[-1] = self.REMOVED
def pop_item(self):
'Remove and return the lowest priority task. Raise KeyError if empty.'
while self.pq:
priority, count, task = heappop(self.pq)
if task is not self.REMOVED:
del self.entry_finder[task]
return task, priority
raise KeyError('pop from an empty priority queue')
def __str__(self):
return str([entry for entry in self.pq if entry[2] != self.REMOVED])
pq = PriorityQueue()
pq.add_task(pq.REMOVED, -100)
pq.add_task(1, -75)
pq.add_task(2, -50)
pq.add_task(pq.REMOVED, -25)
if __name__ == '_main__':
console = []