From d1d7b6a75fd44367c01a284550f65d0e7d6411f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Tue, 28 Nov 2023 11:32:00 +0000
Subject: [PATCH 01/15] Add the heapremove() function

---
 Lib/heapq.py                    | 33 +++++++++++++
 Lib/test/test_heapq.py          | 43 +++++++++++++++++
 Modules/_heapqmodule.c          | 83 +++++++++++++++++++++++++++++++--
 Modules/clinic/_heapqmodule.c.h | 65 +++++++++++++++++++++++++-
 4 files changed, 219 insertions(+), 5 deletions(-)

diff --git a/Lib/heapq.py b/Lib/heapq.py
index 2fd9d1ff4bf827b..ead3bea57beed97 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -167,6 +167,38 @@ def heappushpop(heap, item):
         _siftup(heap, 0)
     return 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.find(item)
+        heapq.heapremove(heap, idx, item)
+    """
+    result = heap[index]
+    if index < 0:
+        index += len(heap)
+    if index < len(heap) - 1:
+        # common case
+        if item is None:
+            item = heap.pop()
+        heap[index] = item
+        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)
+    return result
+
 def heapify(x):
     """Transform list into a heap, in-place, in O(len(x)) time."""
     n = len(x)
@@ -217,6 +249,7 @@ def _siftdown(heap, startpos, pos):
             continue
         break
     heap[pos] = newitem
+    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
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 1aa8e4e289730d0..59beee7d2d4e1e6 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -263,6 +263,49 @@ def __le__(self, other):
         self.assertEqual(hsort(data, LT), target)
         self.assertRaises(TypeError, data, LE)
 
+    def test_remove(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))
+        with self.assertRaises(IndexError) as e:
+            self.module.heapremove(data, -len(data) - 1)
+        
+        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
+            print(len(data), i)
+            v = data[i]
+            self.assertIs(self.module.heapremove(data, i), v)
+            self.check_invariant(data)
+        self.assertFalse(data)
+
+    def test_remove_replace(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))
+        with self.assertRaises(IndexError) as e:
+            self.module.heapremove(data, -len(data) - 1)
+        
+        for i in range(200):
+            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
+            replace = random.random()
+            print(len(data), i, replace)
+            v = data[i]
+            assert self.module.heapremove(data, i, replace) is v
+            self.check_invariant(data)
+
+        self.assertEqual(len(data), 100)
+
 
 class TestHeapPython(TestHeap, TestCase):
     module = py_heapq
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 9d4ec256ee9e3ef..8aac23fa2f64dbe 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -22,10 +22,11 @@ module _heapq
 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=d7cca0a2e4c0ceb3]*/
 
 static int
-siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
+siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t *ppos)
 {
     PyObject *newitem, *parent, **arr;
     Py_ssize_t parentpos, size;
+    Py_ssize_t pos = *ppos;
     int cmp;
 
     assert(PyList_Check(heap));
@@ -63,6 +64,7 @@ siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
         arr[pos] = parent;
         pos = parentpos;
     }
+    *ppos = pos;
     return 0;
 }
 
@@ -113,7 +115,7 @@ siftup(PyListObject *heap, Py_ssize_t pos)
         pos = childpos;
     }
     /* Bubble it up to its final resting place (by sifting its parents down). */
-    return siftdown(heap, startpos, pos);
+    return siftdown(heap, startpos, &pos);
 }
 
 /*[clinic input]
@@ -132,8 +134,8 @@ _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;
 }
@@ -279,6 +281,78 @@ _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item)
     return returnitem;
 }
 
+
+/*[clinic input]
+_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:
+
+    item.value=new_value
+    idx = heap.find(item)
+    heapq.heapremove(heap, idx, item)
+[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=09260f2c67bd4171]*/
+{
+    PyObject *returnitem;
+    Py_ssize_t 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 popped 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);
+        }
+        siftdown((PyListObject *)heap, 0, &index);
+        siftup((PyListObject *)heap, index);
+    } else {
+        /* the tail case */
+        if (item) {
+            PyList_SET_ITEM(heap, index, item);
+            Py_INCREF(item);
+            siftdown((PyListObject *)heap, 0, &index);
+        } else {
+            Py_INCREF(returnitem);
+            if (PyList_SetSlice(heap, index, n, NULL)) {
+                Py_DECREF(returnitem);
+                return NULL;
+            }
+        }
+    }
+    return returnitem;
+}
+
 static Py_ssize_t
 keep_top_bit(Py_ssize_t n)
 {
@@ -540,6 +614,7 @@ static PyMethodDef heapq_methods[] = {
     _HEAPQ__HEAPPOP_MAX_METHODDEF
     _HEAPQ__HEAPIFY_MAX_METHODDEF
     _HEAPQ__HEAPREPLACE_MAX_METHODDEF
+    _HEAPQ_HEAPREMOVE_METHODDEF
     {NULL, NULL}           /* sentinel */
 };
 
diff --git a/Modules/clinic/_heapqmodule.c.h b/Modules/clinic/_heapqmodule.c.h
index 9046307990773ba..10d502d42689efb 100644
--- a/Modules/clinic/_heapqmodule.c.h
+++ b/Modules/clinic/_heapqmodule.c.h
@@ -2,6 +2,7 @@
 preserve
 [clinic start generated code]*/
 
+#include "pycore_abstract.h"      // _PyNumber_Index()
 #include "pycore_modsupport.h"    // _PyArg_CheckPositional()
 
 PyDoc_STRVAR(_heapq_heappush__doc__,
@@ -146,6 +147,68 @@ _heapq_heappushpop(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     return return_value;
 }
 
+PyDoc_STRVAR(_heapq_heapremove__doc__,
+"heapremove($module, heap, index, item=<unrepresentable>, /)\n"
+"--\n"
+"\n"
+"Remove the element at the given index maintaining the heap invariant.\n"
+"\n"
+"An optional item can be provided to replace the removed item. The removed\n"
+"item is returned.\n"
+"This can be used to efficiently remove an item from the heap or\n"
+"to readjust the heap when the comparative \"value\" of\n"
+"an item changes by removing and re-inserting the same item, e.g:\n"
+"\n"
+"    item.value=new_value\n"
+"    idx = heap.find(item)\n"
+"    heapq.heapremove(heap, idx, item)");
+
+#define _HEAPQ_HEAPREMOVE_METHODDEF    \
+    {"heapremove", _PyCFunction_CAST(_heapq_heapremove), METH_FASTCALL, _heapq_heapremove__doc__},
+
+static PyObject *
+_heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index,
+                       PyObject *item);
+
+static PyObject *
+_heapq_heapremove(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    PyObject *heap;
+    Py_ssize_t index;
+    PyObject *item = NULL;
+
+    if (!_PyArg_CheckPositional("heapremove", nargs, 2, 3)) {
+        goto exit;
+    }
+    if (!PyList_Check(args[0])) {
+        _PyArg_BadArgument("heapremove", "argument 1", "list", args[0]);
+        goto exit;
+    }
+    heap = args[0];
+    {
+        Py_ssize_t ival = -1;
+        PyObject *iobj = _PyNumber_Index(args[1]);
+        if (iobj != NULL) {
+            ival = PyLong_AsSsize_t(iobj);
+            Py_DECREF(iobj);
+        }
+        if (ival == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        index = ival;
+    }
+    if (nargs < 3) {
+        goto skip_optional;
+    }
+    item = args[2];
+skip_optional:
+    return_value = _heapq_heapremove_impl(module, heap, index, item);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(_heapq_heapify__doc__,
 "heapify($module, heap, /)\n"
 "--\n"
@@ -267,4 +330,4 @@ _heapq__heapify_max(PyObject *module, PyObject *arg)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=05f2afdf3bc54c9d input=a9049054013a1b77]*/
+/*[clinic end generated code: output=a752a0808801ab63 input=a9049054013a1b77]*/

From 344a35ae1c5e1fe4e467ba4efa50616e7544b8e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Tue, 28 Nov 2023 11:43:34 +0000
Subject: [PATCH 02/15] fixes

---
 Lib/test/test_heapq.py | 10 ++++++++--
 Modules/_heapqmodule.c | 19 +++++++++++--------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 59beee7d2d4e1e6..1686fa7dd027327 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -278,7 +278,7 @@ def test_remove(self):
                 i = 0  # ensure we try 0
             elif random.random() < 0.05:
                 i = len(data) - 1
-            print(len(data), i)
+            # print(len(data), i)
             v = data[i]
             self.assertIs(self.module.heapremove(data, i), v)
             self.check_invariant(data)
@@ -299,13 +299,19 @@ def test_remove_replace(self):
             elif random.random() < 0.05:
                 i = len(data) - 1
             replace = random.random()
-            print(len(data), i, replace)
+            # print(len(data), i, replace)
             v = data[i]
             assert self.module.heapremove(data, i, replace) is v
             self.check_invariant(data)
 
         self.assertEqual(len(data), 100)
 
+    def test_remove_replace_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")
+    
 
 class TestHeapPython(TestHeap, TestCase):
     module = py_heapq
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 8aac23fa2f64dbe..2539e714d42ca23 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -312,20 +312,18 @@ _heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index,
     Py_ssize_t n = PyList_GET_SIZE(heap);
     if (index < 0)
         index += n;
-    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)
-    {
+    if (index < n - 1) {
         /* common case */
         if (item) {
             PyList_SET_ITEM(heap, index, item);
             Py_INCREF(item);
         } else {
-            /* replace with the popped value */
+            /* replace with the last value */
             item = PyList_GET_ITEM(heap, n-1);
             Py_INCREF(item);
             if (PyList_SetSlice(heap, n-1, n, NULL)) {
@@ -334,14 +332,19 @@ _heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index,
             }
             PyList_SET_ITEM(heap, index, item);
         }
-        siftdown((PyListObject *)heap, 0, &index);
-        siftup((PyListObject *)heap, index);
+        if (siftdown((PyListObject *)heap, 0, &index) || siftup((PyListObject *)heap, index)) {
+            Py_DECREF(returnitem);
+            return NULL;
+        }
     } else {
         /* the tail case */
         if (item) {
             PyList_SET_ITEM(heap, index, item);
             Py_INCREF(item);
-            siftdown((PyListObject *)heap, 0, &index);
+            if (siftdown((PyListObject *)heap, 0, &index)) {
+                Py_DECREF(returnitem);
+                return NULL;
+            }
         } else {
             Py_INCREF(returnitem);
             if (PyList_SetSlice(heap, index, n, NULL)) {

From 5b79953a588ebc38ece23d522221b89a68614d78 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Tue, 28 Nov 2023 11:49:54 +0000
Subject: [PATCH 03/15] whitespace

---
 Lib/test/test_heapq.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 1686fa7dd027327..22b22e953dd6546 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -271,7 +271,7 @@ def test_remove(self):
             self.module.heapremove(data, len(data))
         with self.assertRaises(IndexError) as e:
             self.module.heapremove(data, -len(data) - 1)
-        
+
         for i in range(len(data)):
             i = random.randrange(-len(data), len(data))
             if random.random() < 0.05:
@@ -291,7 +291,7 @@ def test_remove_replace(self):
             self.module.heapremove(data, len(data))
         with self.assertRaises(IndexError) as e:
             self.module.heapremove(data, -len(data) - 1)
-        
+
         for i in range(200):
             i = random.randrange(-len(data), len(data))
             if random.random() < 0.05:
@@ -303,7 +303,6 @@ def test_remove_replace(self):
             v = data[i]
             assert self.module.heapremove(data, i, replace) is v
             self.check_invariant(data)
-
         self.assertEqual(len(data), 100)
 
     def test_remove_replace_err(self):
@@ -311,7 +310,7 @@ def test_remove_replace_err(self):
         self.module.heapify(data)
         with self.assertRaisesRegex(TypeError, r"not supported.*'str' and 'float'"):
             self.module.heapremove(data, 50, "foo")
-    
+
 
 class TestHeapPython(TestHeap, TestCase):
     module = py_heapq

From e4ffa87f4aad978e17d698e7a65654ba8d9588fb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Tue, 28 Nov 2023 14:59:32 +0000
Subject: [PATCH 04/15] Add operations count test

---
 Lib/test/test_heapq.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 22b22e953dd6546..9fb6ffc4760614b 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -311,6 +311,47 @@ def test_remove_replace_err(self):
         with self.assertRaisesRegex(TypeError, r"not supported.*'str' and 'float'"):
             self.module.heapremove(data, 50, "foo")
 
+    def test_remove_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(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)
+
 
 class TestHeapPython(TestHeap, TestCase):
     module = py_heapq

From 2af67d70d51bc85d8518b4c9b87b925062930125 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Tue, 28 Nov 2023 17:16:46 +0000
Subject: [PATCH 05/15] documentation

---
 Doc/library/heapq.rst | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index 8b00f7b27959b68..a75d00f743b9342 100644
--- a/Doc/library/heapq.rst
+++ b/Doc/library/heapq.rst
@@ -80,6 +80,21 @@ The following functions are provided:
    on the heap.
 
 
+.. function:: heapremove(heap, index, item=None):
+
+   Remove the item at ``heap[index]``, optionally replacing it with *item*,
+   while maintaining the heap invariant.
+
+   This is more efficient than performing ``del heap[index]`` followed
+   by :func:`heapify`.
+
+   In case an items *value* has changed and its index is known, this
+   function can be used to restore the heap invariant:
+   ``heapremove(heap, item_index, item)``
+
+   .. versionadded:: 3.13
+
+
 The module also offers three general purpose functions based on heaps.
 
 

From 9a8376714c2bfd09bdd66bc1ab6bce02755b22ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Tue, 28 Nov 2023 17:44:38 +0000
Subject: [PATCH 06/15] fix doc

---
 Lib/heapq.py                    | 2 +-
 Modules/_heapqmodule.c          | 4 ++--
 Modules/clinic/_heapqmodule.c.h | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/Lib/heapq.py b/Lib/heapq.py
index ead3bea57beed97..80269b8bbbb3eb7 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -177,7 +177,7 @@ def heapremove(heap, index, item=None):
     an item changes by removing and re-inserting the same item, e.g:
 
         item.value=new_value
-        idx = heap.find(item)
+        idx = heap.index(item)
         heapq.heapremove(heap, idx, item)
     """
     result = heap[index]
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 2539e714d42ca23..56289116ef2de97 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -299,14 +299,14 @@ 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.find(item)
+    idx = heap.index(item)
     heapq.heapremove(heap, idx, item)
 [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=09260f2c67bd4171]*/
+/*[clinic end generated code: output=2d84b49ff0255276 input=13bbd40fb4dee29e]*/
 {
     PyObject *returnitem;
     Py_ssize_t n = PyList_GET_SIZE(heap);
diff --git a/Modules/clinic/_heapqmodule.c.h b/Modules/clinic/_heapqmodule.c.h
index 10d502d42689efb..033af509ab19ce7 100644
--- a/Modules/clinic/_heapqmodule.c.h
+++ b/Modules/clinic/_heapqmodule.c.h
@@ -160,7 +160,7 @@ PyDoc_STRVAR(_heapq_heapremove__doc__,
 "an item changes by removing and re-inserting the same item, e.g:\n"
 "\n"
 "    item.value=new_value\n"
-"    idx = heap.find(item)\n"
+"    idx = heap.index(item)\n"
 "    heapq.heapremove(heap, idx, item)");
 
 #define _HEAPQ_HEAPREMOVE_METHODDEF    \
@@ -330,4 +330,4 @@ _heapq__heapify_max(PyObject *module, PyObject *arg)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=a752a0808801ab63 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=10ce297f0322c817 input=a9049054013a1b77]*/

From c646dc33c9d01fc7433c2966fa0eaf0414b07ae9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Thu, 30 Nov 2023 10:15:38 +0000
Subject: [PATCH 07/15] use heapremove/heapfix as separate functions

---
 Lib/heapq.py                    |  52 +++++++--------
 Lib/test/test_heapq.py          |  15 +++--
 Modules/_heapqmodule.c          | 109 ++++++++++++++------------------
 Modules/clinic/_heapqmodule.c.h |  70 ++++++++++++++------
 4 files changed, 134 insertions(+), 112 deletions(-)

diff --git a/Lib/heapq.py b/Lib/heapq.py
index 80269b8bbbb3eb7..e5fbc02d1858e85 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -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)
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 9fb6ffc4760614b..228209559457ee9 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -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
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 56289116ef2de97..65ed28ad64da525 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -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
index 033af509ab19ce7..e8db9c133c2ee30 100644
--- a/Modules/clinic/_heapqmodule.c.h
+++ b/Modules/clinic/_heapqmodule.c.h
@@ -148,27 +148,18 @@ _heapq_heappushpop(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
 }
 
 PyDoc_STRVAR(_heapq_heapremove__doc__,
-"heapremove($module, heap, index, item=<unrepresentable>, /)\n"
+"heapremove($module, heap, index, /)\n"
 "--\n"
 "\n"
 "Remove the element at the given index maintaining the heap invariant.\n"
 "\n"
-"An optional item can be provided to replace the removed item. The removed\n"
-"item is returned.\n"
-"This can be used to efficiently remove an item from the heap or\n"
-"to readjust the heap when the comparative \"value\" of\n"
-"an item changes by removing and re-inserting the same item, e.g:\n"
-"\n"
-"    item.value=new_value\n"
-"    idx = heap.index(item)\n"
-"    heapq.heapremove(heap, idx, item)");
+"Returns the removed item.");
 
 #define _HEAPQ_HEAPREMOVE_METHODDEF    \
     {"heapremove", _PyCFunction_CAST(_heapq_heapremove), METH_FASTCALL, _heapq_heapremove__doc__},
 
 static PyObject *
-_heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index,
-                       PyObject *item);
+_heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index);
 
 static PyObject *
 _heapq_heapremove(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
@@ -176,9 +167,8 @@ _heapq_heapremove(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     PyObject *return_value = NULL;
     PyObject *heap;
     Py_ssize_t index;
-    PyObject *item = NULL;
 
-    if (!_PyArg_CheckPositional("heapremove", nargs, 2, 3)) {
+    if (!_PyArg_CheckPositional("heapremove", nargs, 2, 2)) {
         goto exit;
     }
     if (!PyList_Check(args[0])) {
@@ -198,12 +188,52 @@ _heapq_heapremove(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
         }
         index = ival;
     }
-    if (nargs < 3) {
-        goto skip_optional;
+    return_value = _heapq_heapremove_impl(module, heap, index);
+
+exit:
+    return return_value;
+}
+
+PyDoc_STRVAR(_heapq_heapfix__doc__,
+"heapfix($module, heap, index, /)\n"
+"--\n"
+"\n"
+"Restore the heap invariant when the element at the given index has been modified.");
+
+#define _HEAPQ_HEAPFIX_METHODDEF    \
+    {"heapfix", _PyCFunction_CAST(_heapq_heapfix), METH_FASTCALL, _heapq_heapfix__doc__},
+
+static PyObject *
+_heapq_heapfix_impl(PyObject *module, PyObject *heap, Py_ssize_t index);
+
+static PyObject *
+_heapq_heapfix(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    PyObject *heap;
+    Py_ssize_t index;
+
+    if (!_PyArg_CheckPositional("heapfix", nargs, 2, 2)) {
+        goto exit;
+    }
+    if (!PyList_Check(args[0])) {
+        _PyArg_BadArgument("heapfix", "argument 1", "list", args[0]);
+        goto exit;
+    }
+    heap = args[0];
+    {
+        Py_ssize_t ival = -1;
+        PyObject *iobj = _PyNumber_Index(args[1]);
+        if (iobj != NULL) {
+            ival = PyLong_AsSsize_t(iobj);
+            Py_DECREF(iobj);
+        }
+        if (ival == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        index = ival;
     }
-    item = args[2];
-skip_optional:
-    return_value = _heapq_heapremove_impl(module, heap, index, item);
+    return_value = _heapq_heapfix_impl(module, heap, index);
 
 exit:
     return return_value;
@@ -330,4 +360,4 @@ _heapq__heapify_max(PyObject *module, PyObject *arg)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=10ce297f0322c817 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=4d721773c4515367 input=a9049054013a1b77]*/

From acd9e06dea298b7b79fe6de715821459fe9a089e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Thu, 30 Nov 2023 11:33:36 +0000
Subject: [PATCH 08/15] docs

---
 Doc/library/heapq.rst | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index a75d00f743b9342..d294270ac6600ef 100644
--- a/Doc/library/heapq.rst
+++ b/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(heap, 0)``.
+
 
 .. function:: heappushpop(heap, item)
 
@@ -80,17 +82,18 @@ The following functions are provided:
    on the heap.
 
 
-.. function:: heapremove(heap, index, item=None):
+.. function:: heapremove(heap, index):
+
+   Remove the item at ``heap[index]`` while maintaining the heap invariant.
+
+   Returns the removed item.
+   
+   .. versionadded:: 3.13
 
-   Remove the item at ``heap[index]``, optionally replacing it with *item*,
-   while maintaining the heap invariant.
 
-   This is more efficient than performing ``del heap[index]`` followed
-   by :func:`heapify`.
+.. function:: heapfix(heap, index):
 
-   In case an items *value* has changed and its index is known, this
-   function can be used to restore the heap invariant:
-   ``heapremove(heap, item_index, item)``
+   Restore the heap invariant after the object at ``heap[index]`` has been modified.
 
    .. versionadded:: 3.13
 

From 5cdb92a31c9b1667fc531fc8568bd404f0a702aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Thu, 30 Nov 2023 11:51:47 +0000
Subject: [PATCH 09/15] whitespace

---
 Doc/library/heapq.rst  | 2 +-
 Lib/heapq.py           | 6 +++---
 Modules/_heapqmodule.c | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index d294270ac6600ef..93e6bf5144f417a 100644
--- a/Doc/library/heapq.rst
+++ b/Doc/library/heapq.rst
@@ -87,7 +87,7 @@ The following functions are provided:
    Remove the item at ``heap[index]`` while maintaining the heap invariant.
 
    Returns the removed item.
-   
+
    .. versionadded:: 3.13
 
 
diff --git a/Lib/heapq.py b/Lib/heapq.py
index e5fbc02d1858e85..ff726a5c1da998f 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -136,7 +136,7 @@ def heappush(heap, item):
 
 def heappop(heap):
     """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
@@ -172,7 +172,7 @@ def heappushpop(heap, item):
 
 def heapremove(heap, index, item=None):
     """Remove the element at the given index maintaining the heap invariant.
-    
+
     Returns the removed object.
     """
     if item is not None:
@@ -184,7 +184,7 @@ def heapremove(heap, index, item=None):
         heap[index] = lastelt
     except IndexError:  # if this was the last item
         return result
-    
+
     if index < 0:
         index += len(heap)
     if index > 0:
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 65ed28ad64da525..5129481088ef29e 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -163,7 +163,7 @@ heapremove_internal(PyObject *heap, Py_ssize_t index, int siftup_func(PyListObje
     }
     if (index == n-1)
         return lastelt;
-    
+
     returnitem = PyList_GET_ITEM(heap, index);
     PyList_SET_ITEM(heap, index, lastelt);
     if (index > 0){

From b297445ab1caaa521860689bba1da39ce6bd7846 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Fri, 1 Dec 2023 09:44:09 +0000
Subject: [PATCH 10/15] fixup

---
 Lib/heapq.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/Lib/heapq.py b/Lib/heapq.py
index ff726a5c1da998f..ca31fd92a3bcd9a 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -170,14 +170,11 @@ def heappushpop(heap, item):
         _siftup(heap, 0)
     return item
 
-def heapremove(heap, index, item=None):
+def heapremove(heap, index):
     """Remove the element at the given index maintaining the heap invariant.
 
     Returns the removed object.
     """
-    if item is not None:
-        heap[index] = item
-        return heapfix(heap, index)
     result = heap[index]  # fails on invalid index
     lastelt = heap.pop()
     try:

From 08cc0f602bbd32a4190b7f8628fa73873e49c18c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Fri, 1 Dec 2023 11:05:06 +0000
Subject: [PATCH 11/15] Fix negative index heapremove

---
 Lib/heapq.py           |  5 +++--
 Lib/test/test_heapq.py | 18 ++++++++++++++++--
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/Lib/heapq.py b/Lib/heapq.py
index ca31fd92a3bcd9a..c1a3023e5dcad54 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -176,14 +176,15 @@ def heapremove(heap, index):
     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
 
-    if index < 0:
-        index += len(heap)
     if index > 0:
         index = _siftdown(heap, 0, index)
     _siftup(heap, index)
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 228209559457ee9..15f10e0bb54c73d 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -266,11 +266,15 @@ def __le__(self, other):
     def test_remove(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(data, len(data))
         with self.assertRaises(IndexError) as e:
             self.module.heapremove(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))
@@ -278,19 +282,26 @@ def test_remove(self):
                 i = 0  # ensure we try 0
             elif random.random() < 0.05:
                 i = len(data) - 1
-            # print(len(data), i)
             v = data[i]
+            # print(len(data), i, v)
+            heapset.remove(v)
             self.assertIs(self.module.heapremove(data, i), v)
+            self.assertEqual(heapset, set(data))
             self.check_invariant(data)
         self.assertFalse(data)
 
     def test_fix(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.heapfix(data, len(data))
         with self.assertRaises(IndexError) as e:
             self.module.heapfix(data, -len(data) - 1)
+        self.check_invariant(data)
+        self.assertEqual(heapset, set(data))
 
         for i in range(200):
             i = random.randrange(-len(data), len(data))
@@ -299,9 +310,12 @@ def test_fix(self):
             elif random.random() < 0.05:
                 i = len(data) - 1
             replace = random.random()
-            # print(len(data), i, replace)
+            # print(len(data), i, data[i], replace)
+            heapset.remove(data[i])
+            heapset.add(replace)
             data[i] = replace
             self.module.heapfix(data, i)
+            self.assertEqual(heapset, set(data))
             self.check_invariant(data)
         self.assertEqual(len(data), 100)
 

From 352764b976afc9a9daff6341bbe2ac9cb8abfca8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Tue, 2 Jan 2024 11:32:23 +0000
Subject: [PATCH 12/15] heapremove now performs a comparison to do either
 "siftup" _or_ "siftdown"

---
 Lib/heapq.py           | 11 ++++++++---
 Modules/_heapqmodule.c | 31 ++++++++++++++++++++++---------
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/Lib/heapq.py b/Lib/heapq.py
index c1a3023e5dcad54..f6dafa95be21733 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -185,9 +185,12 @@ def heapremove(heap, index):
     except IndexError:  # if this was the last item
         return result
 
-    if index > 0:
-        index = _siftdown(heap, 0, index)
-    _siftup(heap, index)
+    # 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 heapfix(heap, index):
@@ -195,6 +198,8 @@ def heapfix(heap, index):
     heap[index]  # trigger an index error for invalid indices
     if index < 0:
         index += len(heap)
+    # we don't have the previous value.  Must sift down first, then up
+    # from the new resting position.
     if index > 0:
         index = _siftdown(heap, 0, index)
     _siftup(heap, index)
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 5129481088ef29e..837b249ef62d97f 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -166,18 +166,27 @@ heapremove_internal(PyObject *heap, Py_ssize_t index, int siftup_func(PyListObje
 
     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){
-        /* 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;
+        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]
@@ -334,6 +343,10 @@ _heapq_heapfix_impl(PyObject *module, PyObject *heap, Py_ssize_t index)
         return NULL;
     }
 
+    /* we cannot compare the value with the previous value so we
+     * must first sift down, and then sift up from the
+     * new resting position
+     */
     if (index > 0){
         if (siftdown((PyListObject *)heap, 0, &index))
             return NULL;

From bbd45a77f566ccfffc2d97d947d15fd19986baad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Wed, 3 Jan 2024 17:54:02 +0000
Subject: [PATCH 13/15] remove heapq.heapfix()

---
 Doc/library/heapq.rst           |  7 -----
 Lib/heapq.py                    | 12 --------
 Lib/test/test_heapq.py          | 36 ------------------------
 Modules/_heapqmodule.c          | 49 +++------------------------------
 Modules/clinic/_heapqmodule.c.h | 47 +------------------------------
 5 files changed, 5 insertions(+), 146 deletions(-)

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index 93e6bf5144f417a..3dc5640f98670b7 100644
--- a/Doc/library/heapq.rst
+++ b/Doc/library/heapq.rst
@@ -91,13 +91,6 @@ The following functions are provided:
    .. versionadded:: 3.13
 
 
-.. function:: heapfix(heap, index):
-
-   Restore the heap invariant after the object at ``heap[index]`` has been modified.
-
-   .. versionadded:: 3.13
-
-
 The module also offers three general purpose functions based on heaps.
 
 
diff --git a/Lib/heapq.py b/Lib/heapq.py
index f6dafa95be21733..c1fe4d9d391aac3 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -193,17 +193,6 @@ def heapremove(heap, 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)
-    # we don't have the previous value.  Must sift down first, then up
-    # from the new resting position.
-    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)
@@ -254,7 +243,6 @@ def _siftdown(heap, startpos, pos):
             continue
         break
     heap[pos] = newitem
-    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
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 15f10e0bb54c73d..4d8a00a6607d60f 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -290,42 +290,6 @@ def test_remove(self):
             self.check_invariant(data)
         self.assertFalse(data)
 
-    def test_fix(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.heapfix(data, len(data))
-        with self.assertRaises(IndexError) as e:
-            self.module.heapfix(data, -len(data) - 1)
-        self.check_invariant(data)
-        self.assertEqual(heapset, set(data))
-
-        for i in range(200):
-            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
-            replace = random.random()
-            # print(len(data), i, data[i], replace)
-            heapset.remove(data[i])
-            heapset.add(replace)
-            data[i] = replace
-            self.module.heapfix(data, i)
-            self.assertEqual(heapset, set(data))
-            self.check_invariant(data)
-        self.assertEqual(len(data), 100)
-
-    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'"):
-            data[50] = "foo"
-            self.module.heapfix(data, 50)
-
     def test_remove_ops(self):
         # test the efficienty of heapremove as opposed to
         # a simple remove and heapify
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 837b249ef62d97f..d1de00440694ee3 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -22,11 +22,10 @@ module _heapq
 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=d7cca0a2e4c0ceb3]*/
 
 static int
-siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t *ppos)
+siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
 {
     PyObject *newitem, *parent, **arr;
     Py_ssize_t parentpos, size;
-    Py_ssize_t pos = *ppos;
     int cmp;
 
     assert(PyList_Check(heap));
@@ -64,7 +63,6 @@ siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t *ppos)
         arr[pos] = parent;
         pos = parentpos;
     }
-    *ppos = pos;
     return 0;
 }
 
@@ -115,7 +113,7 @@ siftup(PyListObject *heap, Py_ssize_t pos)
         pos = childpos;
     }
     /* Bubble it up to its final resting place (by sifting its parents down). */
-    return siftdown(heap, startpos, &pos);
+    return siftdown(heap, startpos, pos);
 }
 
 /*[clinic input]
@@ -135,7 +133,7 @@ _heapq_heappush_impl(PyObject *module, PyObject *heap, PyObject *item)
     if (PyList_Append(heap, item))
         return NULL;
     Py_ssize_t pos = PyList_GET_SIZE(heap)-1;
-    if (siftdown((PyListObject *)heap, 0, &pos))
+    if (siftdown((PyListObject *)heap, 0, pos))
         return NULL;
     Py_RETURN_NONE;
 }
@@ -178,7 +176,7 @@ heapremove_internal(PyObject *heap, Py_ssize_t index, int siftup_func(PyListObje
      * for 'siftdown' which is only called if index > 0
      */
     if (value)
-        value = siftdown((PyListObject *)heap, 0, &index);
+        value = siftdown((PyListObject *)heap, 0, index);
     else
         value = siftup_func((PyListObject *)heap, index);
     if (value)
@@ -318,44 +316,6 @@ _heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index)
     return heapremove_internal(heap, index, siftup);
 }
 
-/*[clinic input]
-_heapq.heapfix
-
-    heap: object(subclass_of='&PyList_Type')
-    index: Py_ssize_t
-    /
-
-Restore the heap invariant when the element at the given index has been modified.
-[clinic start generated code]*/
-
-static PyObject *
-_heapq_heapfix_impl(PyObject *module, PyObject *heap, Py_ssize_t index)
-/*[clinic end generated code: output=c06b38e98a062214 input=2eb44272c71f3106]*/
-{
-    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;
-    }
-
-    /* we cannot compare the value with the previous value so we
-     * must first sift down, and then sift up from the
-     * new resting position
-     */
-    if (index > 0){
-        if (siftdown((PyListObject *)heap, 0, &index))
-            return NULL;
-    }
-    if (siftup((PyListObject *)heap, index))
-        return NULL;
-    Py_RETURN_NONE;
-}
-
 
 static Py_ssize_t
 keep_top_bit(Py_ssize_t n)
@@ -619,7 +579,6 @@ 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
index e8db9c133c2ee30..3a4bdfff2c5ace3 100644
--- a/Modules/clinic/_heapqmodule.c.h
+++ b/Modules/clinic/_heapqmodule.c.h
@@ -194,51 +194,6 @@ _heapq_heapremove(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     return return_value;
 }
 
-PyDoc_STRVAR(_heapq_heapfix__doc__,
-"heapfix($module, heap, index, /)\n"
-"--\n"
-"\n"
-"Restore the heap invariant when the element at the given index has been modified.");
-
-#define _HEAPQ_HEAPFIX_METHODDEF    \
-    {"heapfix", _PyCFunction_CAST(_heapq_heapfix), METH_FASTCALL, _heapq_heapfix__doc__},
-
-static PyObject *
-_heapq_heapfix_impl(PyObject *module, PyObject *heap, Py_ssize_t index);
-
-static PyObject *
-_heapq_heapfix(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
-{
-    PyObject *return_value = NULL;
-    PyObject *heap;
-    Py_ssize_t index;
-
-    if (!_PyArg_CheckPositional("heapfix", nargs, 2, 2)) {
-        goto exit;
-    }
-    if (!PyList_Check(args[0])) {
-        _PyArg_BadArgument("heapfix", "argument 1", "list", args[0]);
-        goto exit;
-    }
-    heap = args[0];
-    {
-        Py_ssize_t ival = -1;
-        PyObject *iobj = _PyNumber_Index(args[1]);
-        if (iobj != NULL) {
-            ival = PyLong_AsSsize_t(iobj);
-            Py_DECREF(iobj);
-        }
-        if (ival == -1 && PyErr_Occurred()) {
-            goto exit;
-        }
-        index = ival;
-    }
-    return_value = _heapq_heapfix_impl(module, heap, index);
-
-exit:
-    return return_value;
-}
-
 PyDoc_STRVAR(_heapq_heapify__doc__,
 "heapify($module, heap, /)\n"
 "--\n"
@@ -360,4 +315,4 @@ _heapq__heapify_max(PyObject *module, PyObject *arg)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=4d721773c4515367 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=2b5eef72a3b55d76 input=a9049054013a1b77]*/

From 11dc5f94fd61d12343fdcfdc70b6362541c66511 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Wed, 3 Jan 2024 18:00:42 +0000
Subject: [PATCH 14/15] Rename heapremove() to heapremove_index()

---
 Doc/library/heapq.rst           |  4 ++--
 Lib/heapq.py                    |  4 ++--
 Lib/test/test_heapq.py          | 12 ++++++------
 Modules/_heapqmodule.c          |  9 +++++----
 Modules/clinic/_heapqmodule.c.h | 21 +++++++++++----------
 5 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index 3dc5640f98670b7..a55849f3bf6bb7d 100644
--- a/Doc/library/heapq.rst
+++ b/Doc/library/heapq.rst
@@ -51,7 +51,7 @@ 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(heap, 0)``.
+   This is equivalent to performing ``heapremove_index(heap, 0)``.
 
 
 .. function:: heappushpop(heap, item)
@@ -82,7 +82,7 @@ The following functions are provided:
    on the heap.
 
 
-.. function:: heapremove(heap, index):
+.. function:: heapremove_index(heap, index):
 
    Remove the item at ``heap[index]`` while maintaining the heap invariant.
 
diff --git a/Lib/heapq.py b/Lib/heapq.py
index c1fe4d9d391aac3..a295c7cb95c5b5e 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -137,7 +137,7 @@ def heappush(heap, item):
 def heappop(heap):
     """Pop the smallest item off the heap, maintaining the heap invariant.
 
-    This is equivalent to heapremove(heap, 0)
+    This is equivalent to heapremove_index(heap, 0)
     """
     lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
     if heap:
@@ -170,7 +170,7 @@ def heappushpop(heap, item):
         _siftup(heap, 0)
     return item
 
-def heapremove(heap, index):
+def heapremove_index(heap, index):
     """Remove the element at the given index maintaining the heap invariant.
 
     Returns the removed object.
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 4d8a00a6607d60f..e34e9e010cc6444 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -263,16 +263,16 @@ def __le__(self, other):
         self.assertEqual(hsort(data, LT), target)
         self.assertRaises(TypeError, data, LE)
 
-    def test_remove(self):
+    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(data, len(data))
+            self.module.heapremove_index(data, len(data))
         with self.assertRaises(IndexError) as e:
-            self.module.heapremove(data, -len(data) - 1)
+            self.module.heapremove_index(data, -len(data) - 1)
         self.check_invariant(data)
         self.assertEqual(heapset, set(data))
 
@@ -285,12 +285,12 @@ def test_remove(self):
             v = data[i]
             # print(len(data), i, v)
             heapset.remove(v)
-            self.assertIs(self.module.heapremove(data, i), 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_ops(self):
+    def test_remove_index_ops(self):
         # test the efficienty of heapremove as opposed to
         # a simple remove and heapify
 
@@ -317,7 +317,7 @@ def __lt__(self, other):
 
         HC.opcount = 0
         while data:
-            self.module.heapremove(data, random.randrange(0, len(data)))
+            self.module.heapremove_index(data, random.randrange(0, len(data)))
         c1 = HC.opcount
 
         data = heapcopy
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index d1de00440694ee3..a6f1d124a5bb276 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -298,7 +298,7 @@ _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item)
 
 
 /*[clinic input]
-_heapq.heapremove
+_heapq.heapremove_index
 
     heap: object(subclass_of='&PyList_Type')
     index: Py_ssize_t
@@ -310,8 +310,9 @@ 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]*/
+_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);
 }
@@ -578,7 +579,7 @@ static PyMethodDef heapq_methods[] = {
     _HEAPQ__HEAPPOP_MAX_METHODDEF
     _HEAPQ__HEAPIFY_MAX_METHODDEF
     _HEAPQ__HEAPREPLACE_MAX_METHODDEF
-    _HEAPQ_HEAPREMOVE_METHODDEF
+    _HEAPQ_HEAPREMOVE_INDEX_METHODDEF
     {NULL, NULL}           /* sentinel */
 };
 
diff --git a/Modules/clinic/_heapqmodule.c.h b/Modules/clinic/_heapqmodule.c.h
index 3a4bdfff2c5ace3..6f94e056ee35e8c 100644
--- a/Modules/clinic/_heapqmodule.c.h
+++ b/Modules/clinic/_heapqmodule.c.h
@@ -147,32 +147,33 @@ _heapq_heappushpop(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     return return_value;
 }
 
-PyDoc_STRVAR(_heapq_heapremove__doc__,
-"heapremove($module, heap, index, /)\n"
+PyDoc_STRVAR(_heapq_heapremove_index__doc__,
+"heapremove_index($module, heap, index, /)\n"
 "--\n"
 "\n"
 "Remove the element at the given index maintaining the heap invariant.\n"
 "\n"
 "Returns the removed item.");
 
-#define _HEAPQ_HEAPREMOVE_METHODDEF    \
-    {"heapremove", _PyCFunction_CAST(_heapq_heapremove), METH_FASTCALL, _heapq_heapremove__doc__},
+#define _HEAPQ_HEAPREMOVE_INDEX_METHODDEF    \
+    {"heapremove_index", _PyCFunction_CAST(_heapq_heapremove_index), METH_FASTCALL, _heapq_heapremove_index__doc__},
 
 static PyObject *
-_heapq_heapremove_impl(PyObject *module, PyObject *heap, Py_ssize_t index);
+_heapq_heapremove_index_impl(PyObject *module, PyObject *heap,
+                             Py_ssize_t index);
 
 static PyObject *
-_heapq_heapremove(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+_heapq_heapremove_index(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
 {
     PyObject *return_value = NULL;
     PyObject *heap;
     Py_ssize_t index;
 
-    if (!_PyArg_CheckPositional("heapremove", nargs, 2, 2)) {
+    if (!_PyArg_CheckPositional("heapremove_index", nargs, 2, 2)) {
         goto exit;
     }
     if (!PyList_Check(args[0])) {
-        _PyArg_BadArgument("heapremove", "argument 1", "list", args[0]);
+        _PyArg_BadArgument("heapremove_index", "argument 1", "list", args[0]);
         goto exit;
     }
     heap = args[0];
@@ -188,7 +189,7 @@ _heapq_heapremove(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
         }
         index = ival;
     }
-    return_value = _heapq_heapremove_impl(module, heap, index);
+    return_value = _heapq_heapremove_index_impl(module, heap, index);
 
 exit:
     return return_value;
@@ -315,4 +316,4 @@ _heapq__heapify_max(PyObject *module, PyObject *arg)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=2b5eef72a3b55d76 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=051f8247c4f85985 input=a9049054013a1b77]*/

From 7dbceb0074706f4fe272856fb3f0d923bc1d22f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Wed, 3 Jan 2024 18:21:09 +0000
Subject: [PATCH 15/15] Add heapremove()

---
 Doc/library/heapq.rst  | 16 ++++++++++++++++
 Lib/heapq.py           | 13 +++++++++++++
 Lib/test/test_heapq.py | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+)

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index a55849f3bf6bb7d..0298df601ccb61a 100644
--- a/Doc/library/heapq.rst
+++ b/Doc/library/heapq.rst
@@ -81,6 +81,22 @@ 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):
 
diff --git a/Lib/heapq.py b/Lib/heapq.py
index a295c7cb95c5b5e..99a44a4ad4d1ff3 100644
--- a/Lib/heapq.py
+++ b/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
@@ -170,6 +172,17 @@ 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.
 
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index e34e9e010cc6444..c7af3da1080f349 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -331,6 +331,42 @@ def __lt__(self, other):
         #  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