add bisect_right and generalize code

benchmark.py

@@ -1,16 +1,16 @@
 from sortedcontainers import SortedList
-from bisect import bisect_left as bisect_left_c
-from bl_bl import bl_bisect_left
+from bisect import bisect_left as bisect_left_c, bisect_right as bisect_right_c
+from bl_bl import bl_bisect_left, bl_bisect_right
 bisect_left = lambda arr, value: arr.bisect_left(value)
 from main import branchless_bisect_left
 from timeit import timeit
 import matplotlib.pyplot as plt
-sizes = [i for i in range(0, 2**20)]
+sizes = [2**i for i in range(0, 25)]
 times = []
 times_me = []
 times_c = []
 times_bl_c = []
-
+# bisect_left
 num_trials = 20
 for size in sizes:
     test_arr = list(range(size))
@@ -31,4 +31,36 @@
 plt.title('Performance vs Input Size')
 plt.legend()
 
+plt.show()
+
+# bisect right
+from bisect import bisect_right as bisect_right_c
+from main import branchless_bisect_right
+from bl_bl import bl_bisect_right
+bisect_right = lambda arr, value: arr.bisect_right(value)
+times = []
+times_me = []
+times_c = []
+times_bl_c = []
+
+num_trials = 20
+for size in sizes:
+    test_arr = list(range(size))
+    value=size//30
+    times_me.append(timeit(lambda:branchless_bisect_right(test_arr, value), number=num_trials))
+    test_arr2 = SortedList(test_arr)
+    times.append(timeit(lambda:bisect_right(test_arr2, value), number=num_trials))
+    times_c.append(timeit(lambda:bisect_right_c(test_arr, value), number=num_trials))
+    times_bl_c.append(timeit(lambda:bl_bisect_right(test_arr, value), number=num_trials))
+    assert len(set([i(test_arr2 if i==bisect_right else test_arr, value) for i in [bisect_right, branchless_bisect_right, bisect_right_c, bl_bisect_right]])) ==1
+
+plt.plot(sizes, times, label="bisect_right")
+plt.plot(sizes, times_c, label="bisect_right_c")
+plt.plot(sizes, times_me, label="branchless_bisect_right")
+plt.plot(sizes, times_bl_c, label="branchless_bisect_right_c")
+plt.xlabel('Input Size')
+plt.ylabel('Execution Time (s)')
+plt.title('Performance vs Input Size')
+plt.legend()
+
 plt.show()

main.py

@@ -13,3 +13,19 @@ def branchless_bisect_left(arr, value):
     for s in (step:=step >> 1 for _ in range(step.bit_length())):
         begin += s*(arr[s+begin]<value)
     return begin+int(arr[begin]<value)
+
+def branchless_bisect_right(arr, value):
+    begin, end=0,len(arr)
+    length = end - begin
+    if length == 0:
+        return end
+    step = 1 << (length.bit_length() - 1)
+    if step != length and arr[step]<=value:
+        length -= step + 1
+        if length == 0:
+            return end
+        begin = end - step
+    step >>= 1
+    for s in (step:=step >> 1 for _ in range(step.bit_length())):
+        begin += s*(arr[s+begin]<=value)
+    return begin+int(arr[begin]<=value)

submodule/main.c

@@ -11,14 +11,7 @@ static long bit_floor_shifted(long x) {
     }
     return res;
 }
-
-static PyObject* bl_bisect_left(PyObject *self, PyObject *args) {
-    PyObject *list_obj;
-    PyObject *value;
-
-    if (!PyArg_ParseTuple(args, "OO", &list_obj, &value))
-        return NULL;
-
+static PyObject* bisect(PyObject *list_obj, PyObject *value, PyObject *compare){
     long size = PyList_Size(list_obj);
 
     long begin = 0, end = size;
@@ -30,25 +23,40 @@ static PyObject* bl_bisect_left(PyObject *self, PyObject *args) {
 
     long step = bit_floor_shifted(length);
 
-    if (step != length && PyObject_RichCompareBool(PyList_GetItem(list_obj, begin + step - 1), value, Py_LT)) {
+    if (step != length && PyObject_RichCompareBool(PyList_GetItem(list_obj, begin + step - 1), value, compare)) {
         begin += step;
         length -= step;
     }
 
     long next = 0;
     for (step >>= 1; step != 0; step >>=1) {
         if ((next = begin + step) < size) {
-            begin += PyObject_RichCompareBool(PyList_GetItem(list_obj, next), value, Py_LT) * step;
+            begin += PyObject_RichCompareBool(PyList_GetItem(list_obj, next), value, compare) * step;
         }
     }
-    return PyLong_FromLong(begin + (begin < size && PyObject_RichCompareBool(PyList_GetItem(list_obj, begin), value, Py_LT)));
+    return PyLong_FromLong(begin + (begin < size && PyObject_RichCompareBool(PyList_GetItem(list_obj, begin), value, compare)));
+}
+static PyObject* bl_bisect_left(PyObject *self, PyObject *args) {
+    PyObject *list_obj;
+    PyObject *value;
+    if (!PyArg_ParseTuple(args, "OO", &list_obj, &value))
+        return NULL;
+    return bisect(list_obj, value, Py_LT);
 }
 
-
+static PyObject* bl_bisect_right(PyObject *self, PyObject *args) {
+    PyObject *list_obj;
+    PyObject *value;
+    if (!PyArg_ParseTuple(args, "OO", &list_obj, &value))
+        return NULL;
+    return bisect(list_obj, value, Py_LE);
+}
 
 static PyMethodDef bl_bl_methods[] = {
     {"bl_bisect_left",  bl_bisect_left, METH_VARARGS,
      "Branchless bisect left"},
+    {"bl_bisect_right",  bl_bisect_right, METH_VARARGS,
+     "Branchless bisect left"},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };