diff --git a/Lib/heapq.py b/Lib/heapq.py
index 0b3e89a3a9..75701c80c5 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -176,6 +176,50 @@ def heapify(x):
for i in reversed(range(n//2)):
_siftup(x, i)
+
+def heapfix(heap, pos):
+ """Re-establish heap invariant after item at index 'pos' has changed.
+
+ Time complexity is O(len(heap)). This is efficient than calling heapify()
+ every time an item changes its value since heapify() performs _siftup()
+ operation on half of the items in the heap, which is unnecessary when we
+ know exactly which element is not at its right place.
+ """
+ # Check if _siftup() succeeds i.e. it moved the item to its correct place
+ # because it was greater than its smaller child (and the child's children
+ # etc). If _siftup() didn't succeed, it means the item *might* move up
+ # using _siftdown(). _siftdown() doesn't move the item up if it is already
+ # in its correct position.
+ if not _siftup(heap, pos):
+ _siftdown(heap, 0, pos)
+
+
+def heapremove(heap, pos):
+ """Remove the item at index 'pos' from the heap.
+
+ Maintains heap invariant after item removal. Time complexity is
+ O(len(heap)).
+
+ Note: Negative 'pos' value is treated as negative index value in list
+ slicing notation i.e. it is counted from the end of the heap.
+ """
+ n = len(heap) - 1
+ # If index position is 0, just use the heappop() method.
+ if pos == 0:
+ return heappop(heap)
+ if pos != n:
+ # Exchange the item at given index with the last element, pop the
+ # exchanged item at the end of the heap and call heapfix() on the
+ # previous last element now at 'pos', which is out of place.
+ heap[pos], heap[n] = heap[n], heap[pos]
+ result = heap.pop()
+ heapfix(heap, pos)
+ else:
+ # If index position is the last index, simply pop that item from the
+ # heap.
+ result = heap.pop()
+ return result
+
def _heappop_max(heap):
"""Maxheap version of a heappop."""
lastelt = heap.pop() # raises appropriate IndexError if heap is empty
@@ -215,6 +259,8 @@ def _siftdown(heap, startpos, pos):
continue
break
heap[pos] = newitem
+ # Return the final position where newitem was inserted.
+ return pos
# The child indices of heap index pos are already heaps, and we want to make
# a heap at index pos too. We do this by bubbling the smaller child of
@@ -257,7 +303,7 @@ def _siftdown(heap, startpos, pos):
def _siftup(heap, pos):
endpos = len(heap)
- startpos = pos
+ startpos = origpos = pos
newitem = heap[pos]
# Bubble up the smaller child until hitting a leaf.
childpos = 2*pos + 1 # leftmost child position
@@ -273,7 +319,12 @@ def _siftup(heap, pos):
# 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)
+ # Get the position where newitem was inserted and return if the new item
+ # moved down the heap from its original (temporary) position because it
+ # was greater than its smaller child (and the child's children etc). This
+ # is determined by checking if newpos > origpos.
+ newpos = _siftdown(heap, startpos, pos)
+ return newpos > origpos
def _siftdown_max(heap, startpos, pos):
'Maxheap variant of _siftdown'