static PyObject * _heapreplace_max(PyObject *self, PyObject *args) { PyObject *heap, *item, *returnitem; if (!PyArg_UnpackTuple(args, "_heapreplace_max", 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 (_siftupmax((PyListObject *)heap, 0) == -1) { Py_DECREF(returnitem); return NULL; } return returnitem; }
int ksmallest(__TYPE *ary, __TYPE len, heap_t *heap) { __TYPE elem, los; __TYPE i, *hp, n; n = heap->tot; heap->ary = ary; hp = ary; heap->len = n; #ifdef ERROR_CHK if(_siftupmax(heap, n/2-1, 0) == -1) return (-1); #else _siftupmax(heap, n/2-1, 0); #endif los = hp[0]; for (i = n; i < len; i++) { elem = ary[i]; if (elem >= los) { continue; } hp[0] = elem; #ifdef ERROR_CHK if (_siftupmax_s(heap, 0) == -1) return (-1); #else _siftupmax_s(heap, 0); #endif los = hp[0]; } return 0; }
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 = PyObject_RichCompareBool(elem, los, Py_LT); 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; Py_DECREF(it); return heap; fail: Py_DECREF(it); Py_XDECREF(heap); return NULL; }