beam: add bisect_bounds (lower/upper bound) with randomized tests

wizard-work/beam/bisect_bounds.py

@@ -0,0 +1,39 @@
+"""Lower/upper bound binary search on a sorted sequence.
+
+Public API:
+    lower_bound(seq, x, *, key=None) -> int
+    upper_bound(seq, x, *, key=None) -> int
+
+lower_bound returns the first index i such that seq[i] >= x (or len(seq) if none).
+upper_bound returns the first index i such that seq[i] > x (or len(seq) if none).
+
+`seq` must be sorted non-decreasingly according to `key` (identity by default).
+"""
+from __future__ import annotations
+from typing import Any, Callable, Sequence
+
+
+def lower_bound(seq: Sequence[Any], x: Any, *, key: Callable[[Any], Any] | None = None) -> int:
+    k = (lambda v: v) if key is None else key
+    target = x
+    lo, hi = 0, len(seq)
+    while lo < hi:
+        mid = (lo + hi) // 2
+        if k(seq[mid]) < target:
+            lo = mid + 1
+        else:
+            hi = mid
+    return lo
+
+
+def upper_bound(seq: Sequence[Any], x: Any, *, key: Callable[[Any], Any] | None = None) -> int:
+    k = (lambda v: v) if key is None else key
+    target = x
+    lo, hi = 0, len(seq)
+    while lo < hi:
+        mid = (lo + hi) // 2
+        if k(seq[mid]) <= target:
+            lo = mid + 1
+        else:
+            hi = mid
+    return lo

wizard-work/beam/test_bisect_bounds.py

@@ -0,0 +1,53 @@
+import os, sys, random
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from bisect_bounds import lower_bound, upper_bound
+
+ok = True
+def check(name, cond):
+    global ok
+    print(("PASS" if cond else "FAIL") + " " + name)
+    ok = ok and cond
+
+# empty
+s = []
+check("empty lower", lower_bound(s, 5) == 0)
+check("empty upper", upper_bound(s, 5) == 0)
+
+# basic
+s = [1,2,3,4,5]
+check("lb 3", lower_bound(s,3)==2)
+check("ub 3", upper_bound(s,3)==3)
+check("lb 0", lower_bound(s,0)==0)
+check("ub 9", upper_bound(s,9)==5)
+check("lb 6", lower_bound(s,6)==5)
+
+# duplicates
+s = [2,2,2,3,3]
+check("lb dup 2", lower_bound(s,2)==0)
+check("ub dup 2", upper_bound(s,2)==3)
+check("count dup 2", upper_bound(s,2)-lower_bound(s,2)==3)
+
+# key
+items = [("a",1),("b",2),("c",2),("d",3)]
+check("key lower", lower_bound(items, 2, key=lambda t: t[1])==1)
+check("key upper", upper_bound(items, 2, key=lambda t: t[1])==3)
+
+# property randomized vs brute force
+random.seed(42)
+for _ in range(2000):
+    n = random.randint(0,30)
+    s = sorted(random.randint(0,9) for _ in range(n))
+    x = random.randint(0,9)
+    lb = lower_bound(s,x)
+    ub = upper_bound(s,x)
+    # verify
+    exp_lb = next((i for i,v in enumerate(s) if v >= x), len(s))
+    exp_ub = next((i for i,v in enumerate(s) if v > x), len(s))
+    if lb != exp_lb or ub != exp_ub:
+        print("MISMATCH", s, x, lb, ub, exp_lb, exp_ub)
+        ok = False
+        break
+check("randomized", ok)
+
+print("ALL" if ok else "SOME FAILED")
+sys.exit(0 if ok else 1)