#!/usr/bin/env python # -*- coding: utf8 -*- """Heap queue algorithm (a.k.a. priority queue).""" # from heapq import heappush, heappop def heappush(heap, item): """Push item onto heap, maintaining the heap invariant.""" heap.append(item) _siftdown(heap, 0, len(heap)-1) def heappop(heap): """Pop the smallest item off the heap, maintaining the heap invariant.""" last_element = heap.pop() # raises appropriate IndexError if heap is empty if heap: returnitem = heap[0] heap[0] = last_element _siftup(heap, 0) return returnitem return last_element def heapreplace(heap, item): """Pop and return the current smallest value, and add the new item. This is more efficient than heappop() followed by heappush(), and can be more appropriate when using a fixed-size heap. Note that the value returned may be larger than item! That constrains reasonable uses of this routine unless written as part of a conditional replacement: if item > heap[0]: item = heapreplace(heap, item) """ returnitem = heap[0] # raises appropriate IndexError if heap is empty heap[0] = item _siftup(heap, 0) return returnitem def heappushpop(heap, item): """Fast version of a heappush followed by a heappop.""" if heap and heap[0] < item: item, heap[0] = heap[0], item _siftup(heap, 0) return item def heapify(x): """Transform list into a heap, in-place, in O(len(x)) time.""" n = len(x) # Transform bottom-up. The largest index there's any point to looking at # is the largest with a child index in-range, so must have 2*i + 1 < n, # or i < (n-1)/2. If n is even = 2*j, this is (2*j-1)/2 = j-1/2 so # j-1 is the largest, which is n//2 - 1. If n is odd = 2*j+1, this is # (2*j+1-1)/2 = j so j-1 is the largest, and that's again n//2-1. for i in reversed(range(n//2)): _siftup(x, i) def _heappop_max(heap): """Maxheap version of a heappop.""" last_element = heap.pop() # raises appropriate IndexError if heap is empty if heap: returnitem = heap[0] heap[0] = last_element _siftup_max(heap, 0) return returnitem return last_element def _heapreplace_max(heap, item): """Maxheap version of a heappop followed by a heappush.""" returnitem = heap[0] # raises appropriate IndexError if heap is empty heap[0] = item _siftup_max(heap, 0) return returnitem def _heapify_max(x): """Transform list into a maxheap, in-place, in O(len(x)) time.""" n = len(x) for i in reversed(range(n//2)): _siftup_max(x, i) def ordered_if_possible(x, y): """Compute x < y, or false if possible (e.g., when pushing on the heap incomparable values like strings and lists).""" try: return x < y except (TypeError, ValueError): return False # 'heap' is a heap at all indices >= startpos, except possibly for pos. pos # is the index of a leaf with a possibly out-of-order value. Restore the # heap invariant. def _siftdown(heap, startpos, pos): newitem = heap[pos] # Follow the path to the root, moving parents down until finding a place # newitem fits. while pos > startpos: parentpos = (pos - 1) >> 1 parent = heap[parentpos] if ordered_if_possible(newitem, parent): heap[pos] = parent pos = parentpos continue break heap[pos] = newitem def _siftup(heap, pos): endpos = len(heap) startpos = pos newitem = heap[pos] # Bubble up the smaller child until hitting a leaf. childpos = 2*pos + 1 # leftmost child position while childpos < endpos: # Set childpos to index of smaller child. rightpos = childpos + 1 if rightpos < endpos and not ordered_if_possible(heap[childpos], heap[rightpos]): childpos = rightpos # Move the smaller child up. heap[pos] = heap[childpos] pos = childpos childpos = 2*pos + 1 # The leaf at pos is empty now. Put newitem there, and bubble it up # to its final resting place (by sifting its parents down). heap[pos] = newitem _siftdown(heap, startpos, pos) def _siftdown_max(heap, startpos, pos): 'Maxheap variant of _siftdown' newitem = heap[pos] # Follow the path to the root, moving parents down until finding a place # newitem fits. while pos > startpos: parentpos = (pos - 1) >> 1 parent = heap[parentpos] if parent < newitem: heap[pos] = parent pos = parentpos continue break heap[pos] = newitem def _siftup_max(heap, pos): 'Maxheap variant of _siftup' endpos = len(heap) startpos = pos newitem = heap[pos] # Bubble up the larger child until hitting a leaf. childpos = 2*pos + 1 # leftmost child position while childpos < endpos: # Set childpos to index of larger child. rightpos = childpos + 1 if rightpos < endpos and not heap[rightpos] < heap[childpos]: childpos = rightpos # Move the larger child up. heap[pos] = heap[childpos] pos = childpos childpos = 2*pos + 1 # The leaf at pos is empty now. Put newitem there, and bubble it up # to its final resting place (by sifting its parents down). heap[pos] = newitem _siftdown_max(heap, startpos, pos) class OurHeap: """ min heap * heap: is the actual heap, heap[1] = index of the smallest element * rank: inverse of heap with rank[x]=i iff heap[i]=x * n: size of the heap :complexity: init O(n log n), len O(1), other operations O(log n) in expectation and O(n) in worst case, due to the usage of a dictionary """ def __init__(self, items): self.heap = [None] # index 0 will be ignored self.rank = {} for x in items: self.push(x) def __len__(self): return len(self.heap) - 1 def push(self, x): """Insert new element x in the heap. Assumption: x is not already in the heap""" assert x not in self.rank i = len(self.heap) self.heap.append(x) # add a new leaf self.rank[x] = i self.up(i) # maintain heap order def pop(self): """Remove and return smallest element""" root = self.heap[1] del self.rank[root] x = self.heap.pop() # remove last leaf if self: # if heap is not empty self.heap[1] = x # put last leaf to root self.rank[x] = 1 self.down(1) # maintain heap order return root def up(self, i): """The value of heap[i] has decreased. Maintain heap invariant.""" x = self.heap[i] while i > 1 and x < self.heap[i // 2]: self.heap[i] = self.heap[i // 2] self.rank[self.heap[i // 2]] = i i //= 2 self.heap[i] = x # insertion index found self.rank[x] = i def down(self, i): """the value of heap[i] has increased. Maintain heap invariant.""" x = self.heap[i] n = len(self.heap) while True: left = 2 * i # climb down the tree right = left + 1 if (right < n and self.heap[right] < x and self.heap[right] < self.heap[left]): self.heap[i] = self.heap[right] self.rank[self.heap[right]] = i # go back up right child i = right elif left < n and self.heap[left] < x: self.heap[i] = self.heap[left] self.rank[self.heap[left]] = i # go back up left child i = left else: self.heap[i] = x # insertion index found self.rank[x] = i return def update(self, old, new): """Replace an element in the heap """ i = self.rank[old] # change value at index i del self.rank[old] self.heap[i] = new self.rank[new] = i if old < new: # maintain heap order self.down(i) else: self.up(i)