/* Drop in replacement for heapq.py  C implementation derived directly from heapq.py in Py2.3which was written by Kevin O'Connor, augmented by Tim Peters,annotated by Fran�ois Pinard, and converted to C by Raymond Hettinger. */ #include "Python.h" static int_siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos){	PyObject *newitem, *parent;	int cmp;	Py_ssize_t 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(parent, newitem, Py_LE);		if (cmp == -1) {			Py_DECREF(newitem);			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_siftup(PyListObject *heap, Py_ssize_t pos){	Py_ssize_t 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, rightpos),				PyList_GET_ITEM(heap, childpos),				Py_LE);			if (cmp == -1) {				Py_DECREF(newitem);				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 _siftdown(heap, startpos, pos);} static PyObject *heappush(PyObject *self, PyObject *args){	PyObject *heap, *item; 	if (!PyArg_UnpackTuple(args, "heappush", 2, 2, &heap, &item))		return NULL; 	if (!PyList_Check(heap)) {		PyErr_SetString(PyExc_TypeError, "heap argument must be a list");		return NULL;	} 	if (PyList_Append(heap, item) == -1)		return NULL; 	if (_siftdown((PyListObject *)heap, 0, PyList_GET_SIZE(heap)-1) == -1)		return NULL;	Py_INCREF(Py_None);	return Py_None;} PyDoc_STRVAR(heappush_doc,"Push item onto heap, maintaining the heap invariant."); static PyObject *heappop(PyObject *self, PyObject *heap){	PyObject *lastelt, *returnitem;	Py_ssize_t n; 	if (!PyList_Check(heap)) {		PyErr_SetString(PyExc_TypeError, "heap argument must be a list");		return NULL;	} 	/* # raises appropriate IndexError if heap is empty */	n = PyList_GET_SIZE(heap);	if (n == 0) {		PyErr_SetString(PyExc_IndexError, "index out of range");		return NULL;	} 	lastelt = PyList_GET_ITEM(heap, n-1) ;	Py_INCREF(lastelt);	PyList_SetSlice(heap, n-1, n, NULL);	n--; 	if (!n) 		return lastelt;	returnitem = PyList_GET_ITEM(heap, 0);	PyList_SET_ITEM(heap, 0, lastelt);	if (_siftup((PyListObject *)heap, 0) == -1) {		Py_DECREF(returnitem);		return NULL;	}	return returnitem;} PyDoc_STRVAR(heappop_doc,"Pop the smallest item off the heap, maintaining the heap invariant."); static PyObject *heapreplace(PyObject *self, PyObject *args){	PyObject *heap, *item, *returnitem; 	if (!PyArg_UnpackTuple(args, "heapreplace", 2, 2, &heap, &item))		return NULL; 	if (!PyList_Check(heap)) {		PyErr_SetString(PyExc_TypeError, "heap argument must be a list");		return NULL;	} 	if (PyList_GET_SIZE(heap) < 1) {		PyErr_SetString(PyExc_IndexError, "index out of range");		return NULL;	} 	returnitem = PyList_GET_ITEM(heap, 0);	Py_INCREF(item);	PyList_SET_ITEM(heap, 0, item);	if (_siftup((PyListObject *)heap, 0) == -1) {		Py_DECREF(returnitem);		return NULL;	}	return returnitem;} PyDoc_STRVAR(heapreplace_doc,"Pop and return the current smallest value, and add the new item.\n\\n\This is more efficient than heappop() followed by heappush(), and can be\n\more appropriate when using a fixed-size heap.  Note that the value\n\returned may be larger than item!  That constrains reasonable uses of\n\this routine unless written as part of a conditional replacement:\n\n\        if item > heap[0]:\n\            item = heapreplace(heap, item)\n"); static PyObject *heapify(PyObject *self, PyObject *heap){	Py_ssize_t i, n; 	if (!PyList_Check(heap)) {		PyErr_SetString(PyExc_TypeError, "heap argument must be a list");		return NULL;	} 	n = PyList_GET_SIZE(heap);	/* Transform bottom-up.  The largest index there's any point to	   looking at is the largest with a child index in-range, so must	   have 2*i + 1 < n, or i < (n-1)/2.  If n is even = 2*j, this is	   (2*j-1)/2 = j-1/2 so j-1 is the largest, which is n//2 - 1.  If	   n is odd = 2*j+1, this is (2*j+1-1)/2 = j so j-1 is the largest,	   and that's again n//2-1.	*/	for (i=n/2-1 ; i>=0 ; i--)		if(_siftup((PyListObject *)heap, i) == -1)			return NULL;	Py_INCREF(Py_None);	return Py_None;} PyDoc_STRVAR(heapify_doc,"Transform list into a heap, in-place, in O(len(heap)) time."); static PyObject *nlargest(PyObject *self, PyObject *args){	PyObject *heap=NULL, *elem, *iterable, *sol, *it, *oldelem;	Py_ssize_t i, n; 	if (!PyArg_ParseTuple(args, "nO:nlargest", &n, &iterable))		return NULL; 	it = PyObject_GetIter(iterable);	if (it == NULL)		return NULL; 	heap = PyList_New(0);	if (heap == NULL)		goto fail; 	for (i=0 ; i<n ; i++ ){		elem = PyIter_Next(it);		if (elem == NULL) {			if (PyErr_Occurred())				goto fail;			else				goto sortit;		}		if (PyList_Append(heap, elem) == -1) {			Py_DECREF(elem);			goto fail;		}		Py_DECREF(elem);	}	if (PyList_GET_SIZE(heap) == 0)		goto sortit; 	for (i=n/2-1 ; i>=0 ; i--)		if(_siftup((PyListObject *)heap, i) == -1)			goto fail; 	sol = 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(elem, sol, Py_LE)) {			Py_DECREF(elem);			continue;		}		oldelem = PyList_GET_ITEM(heap, 0);		PyList_SET_ITEM(heap, 0, elem);		Py_DECREF(oldelem);		if (_siftup((PyListObject *)heap, 0) == -1)			goto fail;		sol = PyList_GET_ITEM(heap, 0);	}sortit:	if (PyList_Sort(heap) == -1)		goto fail;	if (PyList_Reverse(heap) == -1)		goto fail;	Py_DECREF(it);	return heap; fail:	Py_DECREF(it);	Py_XDECREF(heap);	return NULL;} PyDoc_STRVAR(nlargest_doc,"Find the n largest elements in a dataset.\n\\n\Equivalent to:  sorted(iterable, reverse=True)[:n]\n"); static int_siftdownmax(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos){	PyObject *newitem, *parent;	int cmp;	Py_ssize_t 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) {			Py_DECREF(newitem);			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, Py_ssize_t pos){	Py_ssize_t 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) {				Py_DECREF(newitem);				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;	Py_ssize_t i, n; 	if (!PyArg_ParseTuple(args, "nO:nsmallest", &n, &iterable))		return NULL; 	it = PyObject_GetIter(iterable);	if (it == NULL)		return NULL; 	heap = PyList_New(0);	if (heap == NULL)		goto fail; 	for (i=0 ; i<n ; i++ ){		elem = PyIter_Next(it);		if (elem == NULL) {			if (PyErr_Occurred())				goto fail;			else				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:	if (PyList_Sort(heap) == -1)		goto fail;	Py_DECREF(it);	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},	{"heappop",	(PyCFunction)heappop,		METH_O,		heappop_doc},	{"heapreplace",	(PyCFunction)heapreplace,		METH_VARARGS,	heapreplace_doc},	{"heapify",	(PyCFunction)heapify,		METH_O,		heapify_doc},	{"nlargest",	(PyCFunction)nlargest,		METH_VARARGS,	nlargest_doc},	{"nsmallest",	(PyCFunction)nsmallest,		METH_VARARGS,	nsmallest_doc},	{NULL,		NULL}		/* sentinel */}; PyDoc_STRVAR(module_doc,"Heap queue algorithm (a.k.a. priority queue).\n\\n\Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for\n\all k, counting elements from 0.  For the sake of comparison,\n\non-existing elements are considered to be infinite.  The interesting\n\property of a heap is that a[0] is always its smallest element.\n\\n\Usage:\n\\n\heap = []            # creates an empty heap\n\heappush(heap, item) # pushes a new item on the heap\n\item = heappop(heap) # pops the smallest item from the heap\n\item = heap[0]       # smallest item on the heap without popping it\n\heapify(x)           # transforms list into a heap, in-place, in linear time\n\item = heapreplace(heap, item) # pops and returns smallest item, and adds\n\                               # new item; the heap size is unchanged\n\

Lines 1–500 of 619. Load more