Modules/_bisectmodule.c

   1 /* Bisection algorithms. Drop in replacement for bisect.py
   2
   3 Converted to C by Dmitry Vasiliev (dima at hlabs.spb.ru).
   4 */
   5
   6 #include "Python.h"
   7
   8 static Py_ssize_t
   9 internal_bisect_right(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t hi)
  10 {
  11     PyObject *litem;
  12     Py_ssize_t mid, res;
  13
  14     if (lo < 0) {
  15         PyErr_SetString(PyExc_ValueError, "lo must be non-negative");
  16         return -1;
  17     }
  18     if (hi == -1) {
  19         hi = PySequence_Size(list);
  20         if (hi < 0)
  21             return -1;
  22     }
  23     while (lo < hi) {
  24         mid = (lo + hi) / 2;
  25         litem = PySequence_GetItem(list, mid);
  26         if (litem == NULL)
  27             return -1;
  28         res = PyObject_RichCompareBool(item, litem, Py_LT);
  29         Py_DECREF(litem);
  30         if (res < 0)
  31             return -1;
  32         if (res)
  33             hi = mid;
  34         else
  35             lo = mid + 1;
  36     }
  37     return lo;
  38 }
  39
  40 static PyObject *
  41 bisect_right(PyObject *self, PyObject *args, PyObject *kw)
  42 {
  43     PyObject *list, *item;
  44     Py_ssize_t lo = 0;
  45     Py_ssize_t hi = -1;
  46     Py_ssize_t index;
  47     static char *keywords[] = {"a", "x", "lo", "hi", NULL};
  48
  49     if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:bisect_right",
  50         keywords, &list, &item, &lo, &hi))
  51         return NULL;
  52     index = internal_bisect_right(list, item, lo, hi);
  53     if (index < 0)
  54         return NULL;
  55     return PyLong_FromSsize_t(index);
  56 }
  57
  58 PyDoc_STRVAR(bisect_right_doc,
  59 "bisect_right(a, x[, lo[, hi]]) -> index\n\
  60 \n\
  61 Return the index where to insert item x in list a, assuming a is sorted.\n\
  62 \n\
  63 The return value i is such that all e in a[:i] have e <= x, and all e in\n\
  64 a[i:] have e > x.  So if x already appears in the list, i points just\n\
  65 beyond the rightmost x already there\n\
  66 \n\
  67 Optional args lo (default 0) and hi (default len(a)) bound the\n\
  68 slice of a to be searched.\n");
  69
  70 static PyObject *
  71 insort_right(PyObject *self, PyObject *args, PyObject *kw)
  72 {
  73     PyObject *list, *item, *result;
  74     Py_ssize_t lo = 0;
  75     Py_ssize_t hi = -1;
  76     Py_ssize_t index;
  77     static char *keywords[] = {"a", "x", "lo", "hi", NULL};
  78
  79     if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:insort_right",
  80         keywords, &list, &item, &lo, &hi))
  81         return NULL;
  82     index = internal_bisect_right(list, item, lo, hi);
  83     if (index < 0)
  84         return NULL;
  85     if (PyList_CheckExact(list)) {
  86         if (PyList_Insert(list, index, item) < 0)
  87             return NULL;
  88     } else {
  89         result = PyObject_CallMethod(list, "insert", "nO",
  90                                      index, item);
  91         if (result == NULL)
  92             return NULL;
  93         Py_DECREF(result);
  94     }
  95
  96     Py_RETURN_NONE;
  97 }
  98
  99 PyDoc_STRVAR(insort_right_doc,
 100 "insort_right(a, x[, lo[, hi]])\n\
 101 \n\
 102 Insert item x in list a, and keep it sorted assuming a is sorted.\n\
 103 \n\
 104 If x is already in a, insert it to the right of the rightmost x.\n\
 105 \n\
 106 Optional args lo (default 0) and hi (default len(a)) bound the\n\
 107 slice of a to be searched.\n");
 108
 109 static Py_ssize_t
 110 internal_bisect_left(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t hi)
 111 {
 112     PyObject *litem;
 113     Py_ssize_t mid, res;
 114
 115     if (lo < 0) {
 116         PyErr_SetString(PyExc_ValueError, "lo must be non-negative");
 117         return -1;
 118     }
 119     if (hi == -1) {
 120         hi = PySequence_Size(list);
 121         if (hi < 0)
 122             return -1;
 123     }
 124     while (lo < hi) {
 125         mid = (lo + hi) / 2;
 126         litem = PySequence_GetItem(list, mid);
 127         if (litem == NULL)
 128             return -1;
 129         res = PyObject_RichCompareBool(litem, item, Py_LT);
 130         Py_DECREF(litem);
 131         if (res < 0)
 132             return -1;
 133         if (res)
 134             lo = mid + 1;
 135         else
 136             hi = mid;
 137     }
 138     return lo;
 139 }
 140
 141 static PyObject *
 142 bisect_left(PyObject *self, PyObject *args, PyObject *kw)
 143 {
 144     PyObject *list, *item;
 145     Py_ssize_t lo = 0;
 146     Py_ssize_t hi = -1;
 147     Py_ssize_t index;
 148     static char *keywords[] = {"a", "x", "lo", "hi", NULL};
 149
 150     if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:bisect_left",
 151         keywords, &list, &item, &lo, &hi))
 152         return NULL;
 153     index = internal_bisect_left(list, item, lo, hi);
 154     if (index < 0)
 155         return NULL;
 156     return PyLong_FromSsize_t(index);
 157 }
 158
 159 PyDoc_STRVAR(bisect_left_doc,
 160 "bisect_left(a, x[, lo[, hi]]) -> index\n\
 161 \n\
 162 Return the index where to insert item x in list a, assuming a is sorted.\n\
 163 \n\
 164 The return value i is such that all e in a[:i] have e < x, and all e in\n\
 165 a[i:] have e >= x.  So if x already appears in the list, i points just\n\
 166 before the leftmost x already there.\n\
 167 \n\
 168 Optional args lo (default 0) and hi (default len(a)) bound the\n\
 169 slice of a to be searched.\n");
 170
 171 static PyObject *
 172 insort_left(PyObject *self, PyObject *args, PyObject *kw)
 173 {
 174     PyObject *list, *item, *result;
 175     Py_ssize_t lo = 0;
 176     Py_ssize_t hi = -1;
 177     Py_ssize_t index;
 178     static char *keywords[] = {"a", "x", "lo", "hi", NULL};
 179
 180     if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:insort_left",
 181         keywords, &list, &item, &lo, &hi))
 182         return NULL;
 183     index = internal_bisect_left(list, item, lo, hi);
 184     if (index < 0)
 185         return NULL;
 186     if (PyList_CheckExact(list)) {
 187         if (PyList_Insert(list, index, item) < 0)
 188             return NULL;
 189     } else {
 190         result = PyObject_CallMethod(list, "insert", "iO",
 191                                      index, item);
 192         if (result == NULL)
 193             return NULL;
 194         Py_DECREF(result);
 195     }
 196
 197     Py_RETURN_NONE;
 198 }
 199
 200 PyDoc_STRVAR(insort_left_doc,
 201 "insort_left(a, x[, lo[, hi]])\n\
 202 \n\
 203 Insert item x in list a, and keep it sorted assuming a is sorted.\n\
 204 \n\
 205 If x is already in a, insert it to the left of the leftmost x.\n\
 206 \n\
 207 Optional args lo (default 0) and hi (default len(a)) bound the\n\
 208 slice of a to be searched.\n");
 209
 210 PyDoc_STRVAR(bisect_doc, "Alias for bisect_right().\n");
 211 PyDoc_STRVAR(insort_doc, "Alias for insort_right().\n");
 212
 213 static PyMethodDef bisect_methods[] = {
 214     {"bisect_right", (PyCFunction)bisect_right,
 215         METH_VARARGS|METH_KEYWORDS, bisect_right_doc},
 216     {"bisect", (PyCFunction)bisect_right,
 217         METH_VARARGS|METH_KEYWORDS, bisect_doc},
 218     {"insort_right", (PyCFunction)insort_right,
 219         METH_VARARGS|METH_KEYWORDS, insort_right_doc},
 220     {"insort", (PyCFunction)insort_right,
 221         METH_VARARGS|METH_KEYWORDS, insort_doc},
 222     {"bisect_left", (PyCFunction)bisect_left,
 223         METH_VARARGS|METH_KEYWORDS, bisect_left_doc},
 224     {"insort_left", (PyCFunction)insort_left,
 225         METH_VARARGS|METH_KEYWORDS, insort_left_doc},
 226     {NULL, NULL} /* sentinel */
 227 };
 228
 229 PyDoc_STRVAR(module_doc,
 230 "Bisection algorithms.\n\
 231 \n\
 232 This module provides support for maintaining a list in sorted order without\n\
 233 having to sort the list after each insertion. For long lists of items with\n\
 234 expensive comparison operations, this can be an improvement over the more\n\
 235 common approach.\n");
 236
 237
 238 static struct PyModuleDef _bisectmodule = {
 239     PyModuleDef_HEAD_INIT,
 240     "_bisect",
 241     module_doc,
 242     -1,
 243     bisect_methods,
 244     NULL,
 245     NULL,
 246     NULL,
 247     NULL
 248 };
 249
 250 PyMODINIT_FUNC
 251 PyInit__bisect(void)
 252 {
 253     return PyModule_Create(&_bisectmodule);
 254 }