GH-112498 Add heapremove and heapfix

Doc/library/heapq.rst

@@ -51,6 +51,8 @@ The following functions are provided:
    invariant.  If the heap is empty, :exc:`IndexError` is raised.  To access the
    smallest item without popping it, use ``heap[0]``.
 
+   This is equivalent to performing ``heapremove_index(heap, 0)``.
+
 
 .. function:: heappushpop(heap, item)
 
@@ -79,6 +81,31 @@ The following functions are provided:
    combination returns the smaller of the two values, leaving the larger value
    on the heap.
 
+.. function:: heapremove(heap, value, *, key=None):
+
+   Remove the item corresponding to *value* from *heap* while maintaining
+   the heap invariant.
+
+   The heap is searched for the first item comparing equal to *value*.
+
+   If the _callable_ *key* is provided, it will be called for each item
+   before comparing the result with *value*.
+
+   Returns the removed item.
+
+   Raises :exc:`ValueError` if no item corresponding to *value* is found.
+
+   .. versionadded:: 3.13
+
+
+.. function:: heapremove_index(heap, index):
+
+   Remove the item at ``heap[index]`` while maintaining the heap invariant.
+
+   Returns the removed item.
+
+   .. versionadded:: 3.13
+
 
 The module also offers three general purpose functions based on heaps.
 

Lib/heapq.py

@@ -1,3 +1,5 @@
+from operator import indexOf
+
 """Heap queue algorithm (a.k.a. priority queue).
 
 Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for
@@ -135,7 +137,10 @@ def heappush(heap, item):
     _siftdown(heap, 0, len(heap)-1)
 
 def heappop(heap):
-    """Pop the smallest item off the heap, maintaining the heap invariant."""
+    """Pop the smallest item off the heap, maintaining the heap invariant.
+
+    This is equivalent to heapremove_index(heap, 0)
+    """
     lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
     if heap:
         returnitem = heap[0]
@@ -167,6 +172,40 @@ def heappushpop(heap, item):
         _siftup(heap, 0)
     return item
 
+def heapremove(heap, value, *, key = None):
+    """Remove and return the element corresponding to 'value' from the heap, maintaining
+    the heap invariant.
+    The first element comparing equal to 'value' is removed.
+    If a 'key' callable is provided, it is called
+    for each entry in the heap to produce a value to compare with 'value'.
+    Raises ValueError if item is not found in the heap.
+    """
+    index = heap.index(value) if key is None else indexOf(map(key, heap), value)
+    return heapremove_index(heap, index)
+
+def heapremove_index(heap, index):
+    """Remove the element at the given index maintaining the heap invariant.
+
+    Returns the removed object.
+    """
+    result = heap[index]  # fails on invalid index
+    n = len(heap)
+    if index < 0:
+        index += n
+    lastelt = heap.pop()
+    try:
+        heap[index] = lastelt
+    except IndexError:  # if this was the last item
+        return result
+
+    # since we have the old value, we can compare with it and
+    # decide on the direction of sift
+    if index > 0 and lastelt < result:
+        _siftdown(heap, 0, index)
+    else:
+        _siftup(heap, index)
+    return result
+
 def heapify(x):
     """Transform list into a heap, in-place, in O(len(x)) time."""
     n = len(x)

Lib/test/test_heapq.py

@@ -263,6 +263,110 @@ def __le__(self, other):
         self.assertEqual(hsort(data, LT), target)
         self.assertRaises(TypeError, data, LE)
 
+    def test_remove_index(self):
+        data = [random.random() for i in range(100)]
+        self.module.heapify(data)
+        heapset = set(data)
+        self.assertEqual(len(data), len(heapset))
+
+        with self.assertRaises(IndexError) as e:
+            self.module.heapremove_index(data, len(data))
+        with self.assertRaises(IndexError) as e:
+            self.module.heapremove_index(data, -len(data) - 1)
+        self.check_invariant(data)
+        self.assertEqual(heapset, set(data))
+
+        for i in range(len(data)):
+            i = random.randrange(-len(data), len(data))
+            if random.random() < 0.05:
+                i = 0  # ensure we try 0
+            elif random.random() < 0.05:
+                i = len(data) - 1
+            v = data[i]
+            # print(len(data), i, v)
+            heapset.remove(v)
+            self.assertIs(self.module.heapremove_index(data, i), v)
+            self.assertEqual(heapset, set(data))
+            self.check_invariant(data)
+        self.assertFalse(data)
+
+    def test_remove_index_ops(self):
+        # test the efficienty of heapremove as opposed to
+        # a simple remove and heapify
+
+        def dumbremove(heap, index, item=None):
+            if item is None:
+                item = heap.pop()
+                if index == len(heap):
+                    return
+            heap[index] = item
+            self.module.heapify(heap)
+
+        class HC:
+            opcount = 0
+            def __init__(self, val):
+                self.val = val
+            def __lt__(self, other):
+                type(self).opcount += 1
+                return self.val < other.val
+
+        data = [HC(random.random()) for i in range(100)]
+        self.module.heapify(data)
+        heapcopy = data[:]
+        state = random.getstate()
+
+        HC.opcount = 0
+        while data:
+            self.module.heapremove_index(data, random.randrange(0, len(data)))
+        c1 = HC.opcount
+
+        data = heapcopy
+        random.setstate(state)
+        HC.opcount = 0
+        while data:
+            dumbremove(data, random.randrange(0, len(data)))
+        c2 = HC.opcount
+
+        # we should be at least 10 times more efficient (more like 40)
+        #  print(c1, c2)
+        self.assertTrue(c2 > c1 * 10)
+
+    def test_remove(self):
+
+        data = [random.random() for i in range(100)]
+        heap = data[:]
+        self.module.heapify(heap)
+
+        with self.assertRaises(ValueError):
+            self.module.heapremove(heap, "foo")
+
+        item = data[50]
+        self.check_invariant(heap)
+        found = self.module.heapremove(heap, item)
+        self.check_invariant(heap)
+        assert found is item
+        assert found not in heap
+
+    def test_remove_key(self):
+
+        data = [(random.random(), random.random()) for i in range(100)]
+
+        heap = data[:]
+        self.module.heapify(heap)
+
+        key = lambda k: k[1]
+
+        with self.assertRaises(ValueError):
+            self.module.heapremove(heap, "foo", key=key)
+
+        item = data[50]
+        self.check_invariant(heap)
+        assert item in heap
+        found = self.module.heapremove(heap, item[1], key = key)
+        self.check_invariant(heap)
+        assert found is item
+        assert found not in heap
+
 
 class TestHeapPython(TestHeap, TestCase):
     module = py_heapq

Modules/_heapqmodule.c

@@ -132,21 +132,23 @@ _heapq_heappush_impl(PyObject *module, PyObject *heap, PyObject *item)
 {
     if (PyList_Append(heap, item))
         return NULL;
-
-    if (siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1))
+    Py_ssize_t pos = PyList_GET_SIZE(heap)-1;
+    if (siftdown((PyListObject *)heap, 0, pos))
         return NULL;
     Py_RETURN_NONE;
 }
 
 static PyObject *
-heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
+heapremove_internal(PyObject *heap, Py_ssize_t index, int siftup_func(PyListObject *, Py_ssize_t))
 {
     PyObject *lastelt, *returnitem;
     Py_ssize_t n;
 
     /* raises IndexError if the heap is empty */
     n = PyList_GET_SIZE(heap);
-    if (n == 0) {
+    if (index < 0)
+        index += n;
+    if (index < 0 || index >= n) {
         PyErr_SetString(PyExc_IndexError, "index out of range");
         return NULL;
     }
@@ -157,17 +159,32 @@ heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
         Py_DECREF(lastelt);
         return NULL;
     }
-    n--;
-
-    if (!n)
+    if (index == n-1)
         return lastelt;
-    returnitem = PyList_GET_ITEM(heap, 0);
-    PyList_SET_ITEM(heap, 0, lastelt);
-    if (siftup_func((PyListObject *)heap, 0)) {
-        Py_DECREF(returnitem);
-        return NULL;
+
+    returnitem = PyList_GET_ITEM(heap, index);
+    PyList_SET_ITEM(heap, index, lastelt);
+    /* sift down if new item is less than old, else up */
+    int value = 0;
+    if (index > 0){
+        value = PyObject_RichCompareBool(lastelt, returnitem, Py_LT);
+        if (value < 0)
+            goto err;
     }
+    /* the only _max version to call this method is _heappop_max
+     * where index is 0.  So, we need no indirection for the
+     * for 'siftdown' which is only called if index > 0
+     */
+    if (value)
+        value = siftdown((PyListObject *)heap, 0, index);
+    else
+        value = siftup_func((PyListObject *)heap, index);
+    if (value)
+        goto err;
     return returnitem;
+err:
+    Py_DECREF(returnitem);
+    return NULL;
 }
 
 /*[clinic input]
@@ -183,7 +200,7 @@ static PyObject *
 _heapq_heappop_impl(PyObject *module, PyObject *heap)
 /*[clinic end generated code: output=96dfe82d37d9af76 input=91487987a583c856]*/
 {
-    return heappop_internal(heap, siftup);
+    return heapremove_internal(heap, 0, siftup);
 }
 
 static PyObject *
@@ -279,6 +296,28 @@ _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item)
     return returnitem;
 }
 
+
+/*[clinic input]
+_heapq.heapremove_index
+
+    heap: object(subclass_of='&PyList_Type')
+    index: Py_ssize_t
+    /
+
+Remove the element at the given index maintaining the heap invariant.
+
+Returns the removed item.
+[clinic start generated code]*/
+
+static PyObject *
+_heapq_heapremove_index_impl(PyObject *module, PyObject *heap,
+                             Py_ssize_t index)
+/*[clinic end generated code: output=d4157859dfa7484a input=84eff69c62afff1d]*/
+{
+    return heapremove_internal(heap, index, siftup);
+}
+
+
 static Py_ssize_t
 keep_top_bit(Py_ssize_t n)
 {
@@ -494,7 +533,7 @@ static PyObject *
 _heapq__heappop_max_impl(PyObject *module, PyObject *heap)
 /*[clinic end generated code: output=9e77aadd4e6a8760 input=362c06e1c7484793]*/
 {
-    return heappop_internal(heap, siftup_max);
+    return heapremove_internal(heap, 0, siftup_max);
 }
 
 /*[clinic input]
@@ -540,6 +579,7 @@ static PyMethodDef heapq_methods[] = {
     _HEAPQ__HEAPPOP_MAX_METHODDEF
     _HEAPQ__HEAPIFY_MAX_METHODDEF
     _HEAPQ__HEAPREPLACE_MAX_METHODDEF
+    _HEAPQ_HEAPREMOVE_INDEX_METHODDEF
     {NULL, NULL}           /* sentinel */
 };