diff --git a/wizard-work/beam/bisect_bounds.py b/wizard-work/beam/bisect_bounds.py new file mode 100644 index 0000000..5a40f9f --- /dev/null +++ b/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 diff --git a/wizard-work/beam/test_bisect_bounds.py b/wizard-work/beam/test_bisect_bounds.py new file mode 100644 index 0000000..4f1188a --- /dev/null +++ b/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)