diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst index 8b00f7b27959b68..0298df601ccb61a 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_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. diff --git a/Lib/heapq.py b/Lib/heapq.py index 2fd9d1ff4bf827b..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 @@ -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) diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index 1aa8e4e289730d0..c7af3da1080f349 100644 --- a/Lib/test/test_heapq.py +++ b/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 diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c index 9d4ec256ee9e3ef..a6f1d124a5bb276 100644 --- a/Modules/_heapqmodule.c +++ b/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 */ }; diff --git a/Modules/clinic/_heapqmodule.c.h b/Modules/clinic/_heapqmodule.c.h index 9046307990773ba..6f94e056ee35e8c 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,54 @@ _heapq_heappushpop(PyObject *module, PyObject *const *args, Py_ssize_t nargs) return return_value; } +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_INDEX_METHODDEF \ + {"heapremove_index", _PyCFunction_CAST(_heapq_heapremove_index), METH_FASTCALL, _heapq_heapremove_index__doc__}, + +static PyObject * +_heapq_heapremove_index_impl(PyObject *module, PyObject *heap, + Py_ssize_t index); + +static PyObject * +_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_index", nargs, 2, 2)) { + goto exit; + } + if (!PyList_Check(args[0])) { + _PyArg_BadArgument("heapremove_index", "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_heapremove_index_impl(module, heap, index); + +exit: + return return_value; +} + PyDoc_STRVAR(_heapq_heapify__doc__, "heapify($module, heap, /)\n" "--\n" @@ -267,4 +316,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=051f8247c4f85985 input=a9049054013a1b77]*/