use heapremove/heapfix as separate functions

kristjanvalur · Nov 28, 2023 · Nov 28, 2023 · Nov 28, 2023 · Nov 28, 2023 · Nov 28, 2023
commit c646dc33c9d01fc7433c2966fa0eaf0414b07ae9
@@ -135,7 +135,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(heap, 0)
+    """
     lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
     if heap:
         returnitem = heap[0]
@@ -169,36 +172,35 @@ def heappushpop(heap, item):
 
 def heapremove(heap, index, item=None):
     """Remove the element at the given index maintaining the heap invariant.
-
-    An optional item can be provided to replace the removed item. The removed
-    item is returned.
-    This can be used to efficiently remove an item from the heap or
-    to readjust the heap when the comparative "value" of
-    an item changes by removing and re-inserting the same item, e.g:
-
-        item.value=new_value
-        idx = heap.index(item)
-        heapq.heapremove(heap, idx, item)
+
+    Returns the removed object.
     """
-    result = heap[index]
+    if item is not None:
+        heap[index] = item
+        return heapfix(heap, index)
+    result = heap[index]  # fails on invalid index
+    lastelt = heap.pop()
+    try:
+        heap[index] = lastelt
+    except IndexError:  # if this was the last item
+        return result
+
     if index < 0:
         index += len(heap)
-    if index < len(heap) - 1:
-        # common case
-        if item is None:
-            item = heap.pop()
-        heap[index] = item
+    if index > 0:
         index = _siftdown(heap, 0, index)
-        _siftup(heap, index)
-    else:
-        # last item
-        if item is None:
-            heap.pop()
-        else:
-            heap[index] = item
-            _siftdown(heap, 0, index)
+    _siftup(heap, index)
     return result
 
+def heapfix(heap, index):
+    """Restore the heap invariant if element at the given index has changed."""
+    heap[index]  # trigger an index error for invalid indices
+    if index < 0:
+        index += len(heap)
+    if index > 0:
+        index = _siftdown(heap, 0, index)
+    _siftup(heap, index)
+
 def heapify(x):
     """Transform list into a heap, in-place, in O(len(x)) time."""
     n = len(x)

@@ -284,13 +284,13 @@ def test_remove(self):
             self.check_invariant(data)
         self.assertFalse(data)
 
-    def test_remove_replace(self):
+    def test_fix(self):
         data = [random.random() for i in range(100)]
         self.module.heapify(data)
         with self.assertRaises(IndexError) as e:
-            self.module.heapremove(data, len(data))
+            self.module.heapfix(data, len(data))
         with self.assertRaises(IndexError) as e:
-            self.module.heapremove(data, -len(data) - 1)
+            self.module.heapfix(data, -len(data) - 1)
 
         for i in range(200):
             i = random.randrange(-len(data), len(data))
@@ -300,16 +300,17 @@ def test_remove_replace(self):
                 i = len(data) - 1
             replace = random.random()
             # print(len(data), i, replace)
-            v = data[i]
-            assert self.module.heapremove(data, i, replace) is v
+            data[i] = replace
+            self.module.heapfix(data, i)
             self.check_invariant(data)
         self.assertEqual(len(data), 100)
 
-    def test_remove_replace_err(self):
+    def test_fix_err(self):
         data = [random.random() for i in range(100)]
         self.module.heapify(data)
         with self.assertRaisesRegex(TypeError, r"not supported.*'str' and 'float'"):
-            self.module.heapremove(data, 50, "foo")
+            data[50] = "foo"
+            self.module.heapfix(data, 50)
 
     def test_remove_ops(self):
         # test the efficienty of heapremove as opposed to

@@ -141,14 +141,16 @@ _heapq_heappush_impl(PyObject *module, PyObject *heap, PyObject *item)
 }
 
 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;
     }
@@ -159,13 +161,19 @@ 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)) {
+
+    returnitem = PyList_GET_ITEM(heap, index);
+    PyList_SET_ITEM(heap, index, lastelt);
+    if (index > 0){
+        /* there is no max version of heapremove */
+        if (siftdown((PyListObject *)heap, 0, &index)) {
+            Py_DECREF(returnitem);
+            return NULL;
+        }
+    }
+    if (siftup_func((PyListObject *)heap, index)) {
         Py_DECREF(returnitem);
         return NULL;
     }
@@ -185,7 +193,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 *
@@ -287,75 +295,55 @@ _heapq.heapremove
 
     heap: object(subclass_of='&PyList_Type')
     index: Py_ssize_t
-    item: object = NULL
     /
 
 Remove the element at the given index maintaining the heap invariant.
 
-An optional item can be provided to replace the removed item. The removed
-item is returned.
-This can be used to efficiently remove an item from the heap or
-to readjust the heap when the comparative "value" of
-an item changes by removing and re-inserting the same item, e.g:
+Returns the removed item.
+[clinic start generated code]*/
+
+static PyObject *
+_heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index)
+/*[clinic end generated code: output=58466c577a3d8e33 input=b9ab21fc7cfb6e2a]*/
+{
+    return heapremove_internal(heap, index, siftup);
+}
+
+/*[clinic input]
+_heapq.heapfix
+
+    heap: object(subclass_of='&PyList_Type')
+    index: Py_ssize_t
+    /
 
-    item.value=new_value
-    idx = heap.index(item)
-    heapq.heapremove(heap, idx, item)
+Restore the heap invariant when the element at the given index has been modified.
 [clinic start generated code]*/
 
 static PyObject *
-_heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index,
-                       PyObject *item)
-/*[clinic end generated code: output=2d84b49ff0255276 input=13bbd40fb4dee29e]*/
+_heapq_heapfix_impl(PyObject *module, PyObject *heap, Py_ssize_t index)
+/*[clinic end generated code: output=c06b38e98a062214 input=2eb44272c71f3106]*/
 {
-    PyObject *returnitem;
-    Py_ssize_t n = PyList_GET_SIZE(heap);
+    Py_ssize_t n;
+
+    /* raises IndexError if the heap is empty */
+    n = PyList_GET_SIZE(heap);
     if (index < 0)
         index += n;
     if (index < 0 || index >= n) {
         PyErr_SetString(PyExc_IndexError, "index out of range");
         return NULL;
     }
-    returnitem = PyList_GET_ITEM(heap, index);
-    if (index < n - 1) {
-        /* common case */
-        if (item) {
-            PyList_SET_ITEM(heap, index, item);
-            Py_INCREF(item);
-        } else {
-            /* replace with the last value */
-            item = PyList_GET_ITEM(heap, n-1);
-            Py_INCREF(item);
-            if (PyList_SetSlice(heap, n-1, n, NULL)) {
-                Py_DECREF(item);
-                return NULL;
-            }
-            PyList_SET_ITEM(heap, index, item);
-        }
-        if (siftdown((PyListObject *)heap, 0, &index) || siftup((PyListObject *)heap, index)) {
-            Py_DECREF(returnitem);
+
+    if (index > 0){
+        if (siftdown((PyListObject *)heap, 0, &index))
             return NULL;
-        }
-    } else {
-        /* the tail case */
-        if (item) {
-            PyList_SET_ITEM(heap, index, item);
-            Py_INCREF(item);
-            if (siftdown((PyListObject *)heap, 0, &index)) {
-                Py_DECREF(returnitem);
-                return NULL;
-            }
-        } else {
-            Py_INCREF(returnitem);
-            if (PyList_SetSlice(heap, index, n, NULL)) {
-                Py_DECREF(returnitem);
-                return NULL;
-            }
-        }
     }
-    return returnitem;
+    if (siftup((PyListObject *)heap, index))
+        return NULL;
+    Py_RETURN_NONE;
 }
 
+
 static Py_ssize_t
 keep_top_bit(Py_ssize_t n)
 {
@@ -571,7 +559,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]
@@ -618,6 +606,7 @@ static PyMethodDef heapq_methods[] = {
     _HEAPQ__HEAPIFY_MAX_METHODDEF
     _HEAPQ__HEAPREPLACE_MAX_METHODDEF
     _HEAPQ_HEAPREMOVE_METHODDEF
+    _HEAPQ_HEAPFIX_METHODDEF
     {NULL, NULL}           /* sentinel */
 };
 

diff --git a/Modules/clinic/_heapqmodule.c.h b/Modules/clinic/_heapqmodule.c.h