/* Drop in replacement for heapq.py 

C implementation derived directly from heapq.py in Py2.3

which 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"

/* Older implementations of heapq used Py_LE for comparisons.  Now, it uses

   Py_LT so it will match min(), sorted(), and bisect().  Unfortunately, some

   client code (Twisted for example) relied on Py_LE, so this little function

   restores compatability by trying both.

*/

static int

cmp_lt(PyObject *x, PyObject *y)

{

	int cmp;

	static PyObject *lt = NULL;

	if (lt == NULL) {

		lt = PyString_FromString("__lt__");

		if (lt == NULL)

			return -1;

	}

	if (PyObject_HasAttr(x, lt))

		return PyObject_RichCompareBool(x, y, Py_LT);

	cmp = PyObject_RichCompareBool(y, x, Py_LE);

	if (cmp != -1)

		cmp = 1 - cmp;

	return cmp;

}

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 = cmp_lt(newitem, parent);

		if (cmp == -1) {

			Py_DECREF(newitem);

			return -1;

		}

		if (cmp == 0)

			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 = cmp_lt(

				PyList_GET_ITEM(heap, childpos),

				PyList_GET_ITEM(heap, rightpos));

			if (cmp == -1) {

				Py_DECREF(newitem);

				return -1;

			}

			if (cmp == 0)

				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 *

heappushpop(PyObject *self, PyObject *args)

{

	PyObject *heap, *item, *returnitem;

	int cmp;

	if (!PyArg_UnpackTuple(args, "heappushpop", 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) {

		Py_INCREF(item);

		return item;

	}

	cmp = cmp_lt(PyList_GET_ITEM(heap, 0), item);

	if (cmp == -1)

		return NULL;

	if (cmp == 0) {

		Py_INCREF(item);

		return item;

	}

	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(heappushpop_doc,

"Push item on the heap, then pop and return the smallest item\n\

from the heap. The combined action runs more efficiently than\n\

heappush() followed by a separate call to heappop().");

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;

	int cmp;

	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;

		}

		cmp = cmp_lt(sol, elem);

		if (cmp == -1) {

			Py_DECREF(elem);

			goto fail;

		}

		if (cmp == 0) {

			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 = cmp_lt(parent, newitem);

		if (cmp == -1) {

			Py_DECREF(newitem);

			return -1;

		}

		if (cmp == 0)

			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 = cmp_lt(

				PyList_GET_ITEM(heap, rightpos),

				PyList_GET_ITEM(heap, childpos));

			if (cmp == -1) {

				Py_DECREF(newitem);

				return -1;

			}

			if (cmp == 0)

				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;

	int cmp;

	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;

		}

		cmp = cmp_lt(elem, los);

		if (cmp == -1) {

			Py_DECREF(elem);

			goto fail;

		}

		if (cmp == 0) {

			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;