Install C version of heapq.nsmallest().
diff --git a/Doc/lib/libheapq.tex b/Doc/lib/libheapq.tex
index 4585058..5a140ab 100644
--- a/Doc/lib/libheapq.tex
+++ b/Doc/lib/libheapq.tex
@@ -97,11 +97,6 @@
\versionadded{2.4}
\end{funcdesc}
-Though the above functions appear symmetrical, they each have different
-speed and space requirements. In particular, \function{nsmallest()}
-operates on a full copy of the dataset. In contrast, \function{nlargest()}
-only requires storage space for \var{n} elements.
-
Both functions perform best for smaller values of \var{n}. For larger
values, it is more efficient to use the \function{sorted()} function. Also,
when \code{n==1}, it is more efficient to use the builtin \function{min()}
diff --git a/Lib/heapq.py b/Lib/heapq.py
index 65f4155..09f996a 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -300,7 +300,7 @@
# If available, use C implementation
try:
- from _heapq import heappush, heappop, heapify, heapreplace
+ from _heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest
except ImportError:
pass
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 1cdaabe..b6fec9f 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -4,6 +4,7 @@
import random
import unittest
from test import test_support
+import sys
def heapiter(heap):
@@ -91,16 +92,28 @@
def test_nsmallest(self):
data = [random.randrange(2000) for i in range(1000)]
- self.assertEqual(nsmallest(data, 400), sorted(data)[:400])
- self.assertEqual(nsmallest(data, 50), sorted(data)[:50])
+ for i in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
+ self.assertEqual(nsmallest(data, i), sorted(data)[:i])
def test_largest(self):
data = [random.randrange(2000) for i in range(1000)]
- self.assertEqual(nlargest(data, 400), sorted(data, reverse=True)[:400])
+ for i in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
+ self.assertEqual(nlargest(data, i), sorted(data, reverse=True)[:i])
-def test_main():
- test_support.run_unittest(TestHeap)
+def test_main(verbose=None):
+ test_classes = [TestHeap]
+ test_support.run_unittest(*test_classes)
+
+ # verify reference counting
+ if verbose and hasattr(sys, "gettotalrefcount"):
+ import gc
+ counts = [None] * 5
+ for i in xrange(len(counts)):
+ test_support.run_unittest(*test_classes)
+ gc.collect()
+ counts[i] = sys.gettotalrefcount()
+ print counts
if __name__ == "__main__":
- test_main()
+ test_main(verbose=True)
diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 8fe742f..61ee413 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -219,7 +219,7 @@
static PyObject *
nlargest(PyObject *self, PyObject *args)
{
- PyObject *heap=NULL, *elem, *rv, *iterable, *sol, *it, *oldelem;
+ PyObject *heap=NULL, *elem, *iterable, *sol, *it, *oldelem;
int i, n;
if (!PyArg_ParseTuple(args, "Oi:nlargest", &iterable, &n))
@@ -246,10 +246,9 @@
if (PyList_GET_SIZE(heap) == 0)
goto sortit;
- rv = heapify(self, heap);
- if (rv == NULL)
- goto fail;
- Py_DECREF(rv);
+ for (i=n/2-1 ; i>=0 ; i--)
+ if(_siftup((PyListObject *)heap, i) == -1)
+ goto fail;
sol = PyList_GET_ITEM(heap, 0);
while (1) {
@@ -290,6 +289,162 @@
\n\
Equivalent to: sorted(iterable, reverse=True)[:n]\n");
+static int
+_siftdownmax(PyListObject *heap, int startpos, int pos)
+{
+ PyObject *newitem, *parent;
+ int cmp, parentpos;
+
+ assert(PyList_Check(heap));
+ if (pos >= PyList_GET_SIZE(heap)) {
+ PyErr_SetString(PyExc_IndexError, "index out of range");
+ return -1;
+ }
+
+ newitem = PyList_GET_ITEM(heap, pos);
+ Py_INCREF(newitem);
+ /* Follow the path to the root, moving parents down until finding
+ a place newitem fits. */
+ while (pos > startpos){
+ parentpos = (pos - 1) >> 1;
+ parent = PyList_GET_ITEM(heap, parentpos);
+ cmp = PyObject_RichCompareBool(newitem, parent, Py_LE);
+ if (cmp == -1)
+ return -1;
+ if (cmp == 1)
+ break;
+ Py_INCREF(parent);
+ Py_DECREF(PyList_GET_ITEM(heap, pos));
+ PyList_SET_ITEM(heap, pos, parent);
+ pos = parentpos;
+ }
+ Py_DECREF(PyList_GET_ITEM(heap, pos));
+ PyList_SET_ITEM(heap, pos, newitem);
+ return 0;
+}
+
+static int
+_siftupmax(PyListObject *heap, int pos)
+{
+ int startpos, endpos, childpos, rightpos;
+ int cmp;
+ PyObject *newitem, *tmp;
+
+ assert(PyList_Check(heap));
+ endpos = PyList_GET_SIZE(heap);
+ startpos = pos;
+ if (pos >= endpos) {
+ PyErr_SetString(PyExc_IndexError, "index out of range");
+ return -1;
+ }
+ newitem = PyList_GET_ITEM(heap, pos);
+ Py_INCREF(newitem);
+
+ /* Bubble up the smaller child until hitting a leaf. */
+ childpos = 2*pos + 1; /* leftmost child position */
+ while (childpos < endpos) {
+ /* Set childpos to index of smaller child. */
+ rightpos = childpos + 1;
+ if (rightpos < endpos) {
+ cmp = PyObject_RichCompareBool(
+ PyList_GET_ITEM(heap, childpos),
+ PyList_GET_ITEM(heap, rightpos),
+ Py_LE);
+ if (cmp == -1)
+ return -1;
+ if (cmp == 1)
+ childpos = rightpos;
+ }
+ /* Move the smaller child up. */
+ tmp = PyList_GET_ITEM(heap, childpos);
+ Py_INCREF(tmp);
+ Py_DECREF(PyList_GET_ITEM(heap, pos));
+ PyList_SET_ITEM(heap, pos, tmp);
+ pos = childpos;
+ childpos = 2*pos + 1;
+ }
+
+ /* The leaf at pos is empty now. Put newitem there, and and bubble
+ it up to its final resting place (by sifting its parents down). */
+ Py_DECREF(PyList_GET_ITEM(heap, pos));
+ PyList_SET_ITEM(heap, pos, newitem);
+ return _siftdownmax(heap, startpos, pos);
+}
+
+static PyObject *
+nsmallest(PyObject *self, PyObject *args)
+{
+ PyObject *heap=NULL, *elem, *iterable, *los, *it, *oldelem;
+ int i, n;
+
+ if (!PyArg_ParseTuple(args, "Oi:nsmallest", &iterable, &n))
+ return NULL;
+
+ it = PyObject_GetIter(iterable);
+ if (it == NULL)
+ return NULL;
+
+ heap = PyList_New(0);
+ if (it == NULL)
+ goto fail;
+
+ for (i=0 ; i<n ; i++ ){
+ elem = PyIter_Next(it);
+ if (elem == NULL)
+ goto sortit;
+ if (PyList_Append(heap, elem) == -1) {
+ Py_DECREF(elem);
+ goto fail;
+ }
+ Py_DECREF(elem);
+ }
+ n = PyList_GET_SIZE(heap);
+ if (n == 0)
+ goto sortit;
+
+ for (i=n/2-1 ; i>=0 ; i--)
+ if(_siftupmax((PyListObject *)heap, i) == -1)
+ goto fail;
+
+ los = PyList_GET_ITEM(heap, 0);
+ while (1) {
+ elem = PyIter_Next(it);
+ if (elem == NULL) {
+ if (PyErr_Occurred())
+ goto fail;
+ else
+ goto sortit;
+ }
+ if (PyObject_RichCompareBool(los, elem, Py_LE)) {
+ Py_DECREF(elem);
+ continue;
+ }
+
+ oldelem = PyList_GET_ITEM(heap, 0);
+ PyList_SET_ITEM(heap, 0, elem);
+ Py_DECREF(oldelem);
+ if (_siftupmax((PyListObject *)heap, 0) == -1)
+ goto fail;
+ los = PyList_GET_ITEM(heap, 0);
+ }
+
+sortit:
+ Py_DECREF(it);
+ if (PyList_Sort(heap) == -1)
+ goto fail;
+ return heap;
+
+fail:
+ Py_DECREF(it);
+ Py_XDECREF(heap);
+ return NULL;
+}
+
+PyDoc_STRVAR(nsmallest_doc,
+"Find the n smallest elements in a dataset.\n\
+\n\
+Equivalent to: sorted(iterable)[:n]\n");
+
static PyMethodDef heapq_methods[] = {
{"heappush", (PyCFunction)heappush,
METH_VARARGS, heappush_doc},
@@ -301,6 +456,8 @@
METH_O, heapify_doc},
{"nlargest", (PyCFunction)nlargest,
METH_VARARGS, nlargest_doc},
+ {"nsmallest", (PyCFunction)nsmallest,
+ METH_VARARGS, nsmallest_doc},
{NULL, NULL} /* sentinel */
};