Skip to content

Instantly share code, notes, and snippets.

@alanyee
Last active March 3, 2020 04:48
Show Gist options
  • Save alanyee/e0e31f1637f8ad798bf32ce539d1c57e to your computer and use it in GitHub Desktop.
Save alanyee/e0e31f1637f8ad798bf32ce539d1c57e to your computer and use it in GitHub Desktop.
Maxheap version of heapq.py for safe keeping
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)
def heappop_max(heap):
"""Maxheap version of a heappop."""
lastelt = heap.pop() # raises appropriate IndexError if heap is empty
if heap:
returnitem = heap[0]
heap[0] = lastelt
_siftup_max(heap, 0)
return returnitem
return lastelt
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 heappush_max(heap, item):
"""Maxheap version of heappush."""
heap.append(item)
_siftdown_max(heap, 0, len(heap)-1)
def heappushpop_max(heap, item):
"""Maxheap version of heappushpop."""
if heap and heap[0] > item:
item, heap[0] = heap[0], item
_siftup_max(heap, 0)
return item
# If available, use C implementation
try:
from _heapq import _heapreplace_max as heapreplace_max
except ImportError:
pass
try:
from _heapq import _heapify_max as heapify_max
except ImportError:
pass
try:
from _heapq import _heappop_max as heappop_max
except ImportError:
pass
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment