40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
"""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
|