Binary Search Bounds With Keys

from __future__ import annotations

from collections.abc import Callable, Sequence
from typing import TypeVar

T = TypeVar("T")
K = TypeVar("K")


def lower_bound(
    items: Sequence[T],
    target: K,
    *,
    key: Callable[[T], K] | None = None,
) -> int:
    low, high = 0, len(items)
    while low < high:
        middle = (low + high) // 2
        value = key(items[middle]) if key is not None else items[middle]
        if value < target:
            low = middle + 1
        else:
            high = middle
    return low


def upper_bound(
    items: Sequence[T],
    target: K,
    *,
    key: Callable[[T], K] | None = None,
) -> int:
    low, high = 0, len(items)
    while low < high:
        middle = (low + high) // 2
        value = key(items[middle]) if key is not None else items[middle]
        if value <= target:
            low = middle + 1
        else:
            high = middle
    return low


def binary_search(
    items: Sequence[T],
    target: K,
    *,
    key: Callable[[T], K] | None = None,
) -> int:
    index = lower_bound(items, target, key=key)
    if index >= len(items):
        return -1
    value = key(items[index]) if key is not None else items[index]
    return index if value == target else -1


def count_equal(
    items: Sequence[T],
    target: K,
    *,
    key: Callable[[T], K] | None = None,
) -> int:
    return upper_bound(items, target, key=key) - lower_bound(
        items,
        target,
        key=key,
    )


def main() -> None:
    numbers = [1, 3, 3, 5, 8]
    print(lower_bound(numbers, 3), upper_bound(numbers, 3))
    print(count_equal(numbers, 3))
    print(binary_search(numbers, 5), binary_search(numbers, 4))

    records = [
        {"name": "ada", "age": 31},
        {"name": "grace", "age": 36},
        {"name": "linus", "age": 55},
    ]
    print(lower_bound(records, 36, key=lambda record: record["age"]))
    print(
        binary_search(
            records,
            55,
            key=lambda record: record["age"],
        )
    )

    words = ["apple", "banana", "cherry"]
    print(lower_bound(words, "blueberry"))


if __name__ == "__main__":
    main()