Skip to content

bpo-37229: Add compare_function to bisect functions #13970

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

bpo-37229: Add compare_function to bisect functions #13970

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 16 additions & 8 deletions Lib/bisect.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,20 @@
"""Bisection algorithms."""

def insort_right(a, x, lo=0, hi=None):
def insort_right(a, x, lo=0, hi=None, compare_function=None):
"""Insert item x in list a, and keep it sorted assuming a is sorted.

If x is already in a, insert it to the right of the rightmost x.

Optional args lo (default 0) and hi (default len(a)) bound the
slice of a to be searched.
By default, comparison is done via __lt__ (such as a < b).
compare_function can be supplied to provide an alternative boolean comparison of the same order.
"""

lo = bisect_right(a, x, lo, hi)
lo = bisect_right(a, x, lo, hi, compare_function)
a.insert(lo, x)

def bisect_right(a, x, lo=0, hi=None):
def bisect_right(a, x, lo=0, hi=None, compare_function=None):
"""Return the index where to insert item x in list a, assuming a is sorted.

The return value i is such that all e in a[:i] have e <= x, and all e in
Expand All @@ -21,6 +23,8 @@ def bisect_right(a, x, lo=0, hi=None):

Optional args lo (default 0) and hi (default len(a)) bound the
slice of a to be searched.
By default, comparison is done via __lt__ (such as a < b).
compare_function can be supplied to provide an alternative boolean comparison of the same order.
"""

if lo < 0:
Expand All @@ -29,24 +33,26 @@ def bisect_right(a, x, lo=0, hi=None):
hi = len(a)
while lo < hi:
mid = (lo+hi)//2
if x < a[mid]: hi = mid
if x < a[mid] if compare_function is None else compare_function(x, a[mid]): hi = mid
else: lo = mid+1
return lo

def insort_left(a, x, lo=0, hi=None):
def insort_left(a, x, lo=0, hi=None, compare_function=None):
"""Insert item x in list a, and keep it sorted assuming a is sorted.

If x is already in a, insert it to the left of the leftmost x.

Optional args lo (default 0) and hi (default len(a)) bound the
slice of a to be searched.
By default, comparison is done via __lt__ (such as a < b).
compare_function can be supplied to provide an alternative boolean comparison of the same order.
"""

lo = bisect_left(a, x, lo, hi)
lo = bisect_left(a, x, lo, hi, compare_function)
a.insert(lo, x)


def bisect_left(a, x, lo=0, hi=None):
def bisect_left(a, x, lo=0, hi=None, compare_function=None):
"""Return the index where to insert item x in list a, assuming a is sorted.

The return value i is such that all e in a[:i] have e < x, and all e in
Expand All @@ -55,6 +61,8 @@ def bisect_left(a, x, lo=0, hi=None):

Optional args lo (default 0) and hi (default len(a)) bound the
slice of a to be searched.
By default, comparison is done via __lt__ (such as a < b).
compare_function can be supplied to provide an alternative boolean comparison of the same order.
"""

if lo < 0:
Expand All @@ -63,7 +71,7 @@ def bisect_left(a, x, lo=0, hi=None):
hi = len(a)
while lo < hi:
mid = (lo+hi)//2
if a[mid] < x: lo = mid+1
if a[mid] < x if compare_function is None else compare_function(a[mid], x): lo = mid+1
else: hi = mid
return lo

Expand Down
20 changes: 20 additions & 0 deletions Lib/test/test_bisect.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,26 @@ def test_keyword_args(self):
self.module.insort(a=data, x=25, lo=1, hi=3)
self.assertEqual(data, [10, 20, 25, 25, 25, 30, 40, 50])

def test_compare_function(self):
def compare_function(c1, c2):
return c1.n < c2.n

class CustomClass:
def __init__(self, n):
self.n = n

def __eq__(self, other):
return self.n == other.n

data = [CustomClass(n) for n in (10, 20, 30, 40, 50)]
self.assertEqual(self.module.bisect_left(a=data, x=CustomClass(25), lo=1, hi=3, compare_function=compare_function), 2)
self.assertEqual(self.module.bisect_right(a=data, x=CustomClass(25), lo=1, hi=3, compare_function=compare_function), 2)
self.assertEqual(self.module.bisect(a=data, x=CustomClass(25), lo=1, hi=3, compare_function=compare_function), 2)
self.module.insort_left(a=data, x=CustomClass(25), lo=1, hi=3, compare_function=compare_function)
self.module.insort_right(a=data, x=CustomClass(25), lo=1, hi=3, compare_function=compare_function)
self.module.insort(a=data, x=CustomClass(25), lo=1, hi=3, compare_function=compare_function)
self.assertEqual(data, [CustomClass(10), CustomClass(20), CustomClass(25), CustomClass(25), CustomClass(25), CustomClass(30), CustomClass(40), CustomClass(50)])

class TestBisectPython(TestBisect, unittest.TestCase):
module = py_bisect

Expand Down
96 changes: 64 additions & 32 deletions Modules/_bisectmodule.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,12 @@ Converted to C by Dmitry Vasiliev (dima at hlabs.spb.ru).
_Py_IDENTIFIER(insert);

static inline Py_ssize_t
internal_bisect_right(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t hi)
internal_bisect_right(PyObject *list, PyObject *item, PyObject *func, Py_ssize_t lo, Py_ssize_t hi)
{
PyObject *litem;
Py_ssize_t mid;
int res;
int res = -1;
PyObject *callback_res;

if (lo < 0) {
PyErr_SetString(PyExc_ValueError, "lo must be non-negative");
Expand All @@ -32,7 +33,16 @@ internal_bisect_right(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t
litem = PySequence_GetItem(list, mid);
if (litem == NULL)
return -1;
res = PyObject_RichCompareBool(item, litem, Py_LT);
if (NULL != func && NULL != func->ob_type && NULL!= func->ob_type->tp_call && PyCallable_Check(func)) {
callback_res = PyObject_CallFunctionObjArgs(func, item, litem, NULL);
if (NULL != callback_res) {
PyArg_Parse(callback_res, "p", &res);
} else {
res = -1;
}
} else {
res = PyObject_RichCompareBool(item, litem, Py_LT);
}
Py_DECREF(litem);
if (res < 0)
return -1;
Expand All @@ -47,29 +57,29 @@ internal_bisect_right(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t
static PyObject *
bisect_right(PyObject *self, PyObject *args, PyObject *kw)
{
PyObject *list, *item;
PyObject *list, *item, *func = NULL;
Py_ssize_t lo = 0;
Py_ssize_t hi = -1;
Py_ssize_t index;
static char *keywords[] = {"a", "x", "lo", "hi", NULL};
static char *keywords[] = {"a", "x", "lo", "hi", "compare_function", NULL};

if (kw == NULL && PyTuple_GET_SIZE(args) == 2) {
list = PyTuple_GET_ITEM(args, 0);
item = PyTuple_GET_ITEM(args, 1);
}
else {
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:bisect_right",
keywords, &list, &item, &lo, &hi))
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nnO:bisect_right",
keywords, &list, &item, &lo, &hi, &func))
return NULL;
}
index = internal_bisect_right(list, item, lo, hi);
index = internal_bisect_right(list, item, func, lo, hi);
if (index < 0)
return NULL;
return PyLong_FromSsize_t(index);
}

PyDoc_STRVAR(bisect_right_doc,
"bisect_right(a, x[, lo[, hi]]) -> index\n\
"bisect_right(a, x[, lo[, hi[, compare_function]]]) -> index\n\
\n\
Return the index where to insert item x in list a, assuming a is sorted.\n\
\n\
Expand All @@ -78,27 +88,30 @@ a[i:] have e > x. So if x already appears in the list, i points just\n\
beyond the rightmost x already there\n\
\n\
Optional args lo (default 0) and hi (default len(a)) bound the\n\
slice of a to be searched.\n");
slice of a to be searched.\n\
By default, comparison is done via __lt__ (such as a < b).\
\n\
compare_function can be supplied to provide an alternative boolean comparison of the same order.");

static PyObject *
insort_right(PyObject *self, PyObject *args, PyObject *kw)
{
PyObject *list, *item, *result;
PyObject *list, *item, *result, *func = NULL;
Py_ssize_t lo = 0;
Py_ssize_t hi = -1;
Py_ssize_t index;
static char *keywords[] = {"a", "x", "lo", "hi", NULL};
static char *keywords[] = {"a", "x", "lo", "hi", "compare_function", NULL};

if (kw == NULL && PyTuple_GET_SIZE(args) == 2) {
list = PyTuple_GET_ITEM(args, 0);
item = PyTuple_GET_ITEM(args, 1);
}
else {
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:insort_right",
keywords, &list, &item, &lo, &hi))
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nnO:insort_right",
keywords, &list, &item, &lo, &hi, &func))
return NULL;
}
index = internal_bisect_right(list, item, lo, hi);
index = internal_bisect_right(list, item, func, lo, hi);
if (index < 0)
return NULL;
if (PyList_CheckExact(list)) {
Expand All @@ -116,21 +129,25 @@ insort_right(PyObject *self, PyObject *args, PyObject *kw)
}

PyDoc_STRVAR(insort_right_doc,
"insort_right(a, x[, lo[, hi]])\n\
"insort_right(a, x[, lo[, hi, [compare_function]]])\n\
\n\
Insert item x in list a, and keep it sorted assuming a is sorted.\n\
\n\
If x is already in a, insert it to the right of the rightmost x.\n\
\n\
Optional args lo (default 0) and hi (default len(a)) bound the\n\
slice of a to be searched.\n");
slice of a to be searched.\n\
By default, comparison is done via __lt__ (such as a < b).\
\n\
compare_function can be supplied to provide an alternative boolean comparison of the same order.");

static inline Py_ssize_t
internal_bisect_left(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t hi)
internal_bisect_left(PyObject *list, PyObject *item, PyObject *func, Py_ssize_t lo, Py_ssize_t hi)
{
PyObject *litem;
Py_ssize_t mid;
int res;
int res = -1;
PyObject *callback_res;

if (lo < 0) {
PyErr_SetString(PyExc_ValueError, "lo must be non-negative");
Expand All @@ -149,7 +166,16 @@ internal_bisect_left(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t h
litem = PySequence_GetItem(list, mid);
if (litem == NULL)
return -1;
res = PyObject_RichCompareBool(litem, item, Py_LT);
if (NULL != func && NULL != func->ob_type && NULL != func->ob_type->tp_call && PyCallable_Check(func)) {
callback_res = PyObject_CallFunctionObjArgs(func, litem, item, NULL);
if (NULL != callback_res) {
PyArg_Parse(callback_res, "p", &res);
} else {
res = -1;
}
} else {
res = PyObject_RichCompareBool(litem, item, Py_LT);
}
Py_DECREF(litem);
if (res < 0)
return -1;
Expand All @@ -164,22 +190,22 @@ internal_bisect_left(PyObject *list, PyObject *item, Py_ssize_t lo, Py_ssize_t h
static PyObject *
bisect_left(PyObject *self, PyObject *args, PyObject *kw)
{
PyObject *list, *item;
PyObject *list, *item, *func = NULL;
Py_ssize_t lo = 0;
Py_ssize_t hi = -1;
Py_ssize_t index;
static char *keywords[] = {"a", "x", "lo", "hi", NULL};
static char *keywords[] = {"a", "x", "lo", "hi", "compare_function", NULL};

if (kw == NULL && PyTuple_GET_SIZE(args) == 2) {
list = PyTuple_GET_ITEM(args, 0);
item = PyTuple_GET_ITEM(args, 1);
}
else {
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:bisect_left",
keywords, &list, &item, &lo, &hi))
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nnO:bisect_left",
keywords, &list, &item, &lo, &hi, &func))
return NULL;
}
index = internal_bisect_left(list, item, lo, hi);
index = internal_bisect_left(list, item, func, lo, hi);
if (index < 0)
return NULL;
return PyLong_FromSsize_t(index);
Expand All @@ -195,26 +221,29 @@ a[i:] have e >= x. So if x already appears in the list, i points just\n\
before the leftmost x already there.\n\
\n\
Optional args lo (default 0) and hi (default len(a)) bound the\n\
slice of a to be searched.\n");
slice of a to be searched.\n\
By default, comparison is done via __lt__ (such as a < b).\
\n\
compare_function can be supplied to provide an alternative boolean comparison of the same order.");

static PyObject *
insort_left(PyObject *self, PyObject *args, PyObject *kw)
{
PyObject *list, *item, *result;
PyObject *list, *item, *result, *func = NULL;
Py_ssize_t lo = 0;
Py_ssize_t hi = -1;
Py_ssize_t index;
static char *keywords[] = {"a", "x", "lo", "hi", NULL};
static char *keywords[] = {"a", "x", "lo", "hi", "compare_function", NULL};

if (kw == NULL && PyTuple_GET_SIZE(args) == 2) {
list = PyTuple_GET_ITEM(args, 0);
item = PyTuple_GET_ITEM(args, 1);
} else {
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nn:insort_left",
keywords, &list, &item, &lo, &hi))
if (!PyArg_ParseTupleAndKeywords(args, kw, "OO|nnO:insort_left",
keywords, &list, &item, &lo, &hi, &func))
return NULL;
}
index = internal_bisect_left(list, item, lo, hi);
index = internal_bisect_left(list, item, func, lo, hi);
if (index < 0)
return NULL;
if (PyList_CheckExact(list)) {
Expand All @@ -238,7 +267,10 @@ Insert item x in list a, and keep it sorted assuming a is sorted.\n\
If x is already in a, insert it to the left of the leftmost x.\n\
\n\
Optional args lo (default 0) and hi (default len(a)) bound the\n\
slice of a to be searched.\n");
slice of a to be searched.\n\
By default, comparison is done via __lt__ (such as a < b).\
\n\
compare_function can be supplied to provide an alternative boolean comparison of the same order.");

static PyMethodDef bisect_methods[] = {
{"bisect_right", (PyCFunction)(void(*)(void))bisect_right,
Expand Down