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 PyInt_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 PyInt_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 PyMODINIT_FUNC
 238 init_bisect(void)
 239 {
 240         PyObject *m;
 241
 242         m = Py_InitModule3("_bisect", bisect_methods, module_doc);
 243 }