| 55 | |
| 56 | |
| 57 | class PriorityQueue: |
| 58 | def __init__(self, capacity, heap_order="max"): |
| 59 | """ |
| 60 | A priority queue implementation using a binary heap. |
| 61 | |
| 62 | Notes |
| 63 | ----- |
| 64 | A priority queue is a data structure useful for storing the top |
| 65 | `capacity` largest or smallest elements in a collection of values. As a |
| 66 | result of using a binary heap, ``PriorityQueue`` offers `O(log N)` |
| 67 | :meth:`push` and :meth:`pop` operations. |
| 68 | |
| 69 | Parameters |
| 70 | ---------- |
| 71 | capacity: int |
| 72 | The maximum number of items that can be held in the queue. |
| 73 | heap_order: {"max", "min"} |
| 74 | Whether the priority queue should retain the items with the |
| 75 | `capacity` smallest (`heap_order` = 'min') or `capacity` largest |
| 76 | (`heap_order` = 'max') priorities. |
| 77 | """ |
| 78 | assert heap_order in ["max", "min"], "heap_order must be either 'max' or 'min'" |
| 79 | self.capacity = capacity |
| 80 | self.heap_order = heap_order |
| 81 | |
| 82 | self._pq = [] |
| 83 | self._count = 0 |
| 84 | self._entry_counter = 0 |
| 85 | |
| 86 | def __repr__(self): |
| 87 | fstr = "PriorityQueue(capacity={}, heap_order={}) with {} items" |
| 88 | return fstr.format(self.capacity, self.heap_order, self._count) |
| 89 | |
| 90 | def __len__(self): |
| 91 | return self._count |
| 92 | |
| 93 | def __iter__(self): |
| 94 | return iter(self._pq) |
| 95 | |
| 96 | def push(self, key, priority, val=None): |
| 97 | """ |
| 98 | Add a new (key, value) pair with priority `priority` to the queue. |
| 99 | |
| 100 | Notes |
| 101 | ----- |
| 102 | If the queue is at capacity and `priority` exceeds the priority of the |
| 103 | item with the largest/smallest priority currently in the queue, replace |
| 104 | the current queue item with (`key`, `val`). |
| 105 | |
| 106 | Parameters |
| 107 | ---------- |
| 108 | key : hashable object |
| 109 | The key to insert into the queue. |
| 110 | priority : comparable |
| 111 | The priority for the `key`, `val` pair. |
| 112 | val : object |
| 113 | The value associated with `key`. Default is None. |
| 114 | """ |