"""Binary search, partition, and maximum subarray with invariants. Python 3.10+."""
from __future__ import annotations
from bisect import bisect_left
from itertools import product


def lower_bound(values: list[int], target: int) -> int:
    """Pre: ascending integers. Post: first index with value>=target, or len(values)."""
    if any(a > b for a, b in zip(values, values[1:])):
        raise ValueError("values must be sorted")
    lo, hi = 0, len(values)
    # Invariant: values[:lo]<target, values[hi:]>=target, 0<=lo<=hi<=n.
    # Bound: hi-lo.
    while lo < hi:
        mid = (lo + hi) // 2
        if values[mid] < target:
            lo = mid + 1
        else:
            hi = mid
    return lo


def sort_colors(values: list[int]) -> list[int]:
    """Pre: only 0,1,2. Post: a sorted copy, preserving the caller's list."""
    if any(v not in (0, 1, 2) for v in values):
        raise ValueError("colors must be 0, 1, or 2")
    result = list(values)
    lo, i, hi = 0, 0, len(result)
    # Invariant: zeros before lo, ones in [lo,i), twos from hi; permutation preserved.
    # Bound: hi-i, the unknown region.
    while i < hi:
        if result[i] == 0:
            result[lo], result[i] = result[i], result[lo]
            lo, i = lo + 1, i + 1
        elif result[i] == 1:
            i += 1
        else:
            hi -= 1
            result[i], result[hi] = result[hi], result[i]
    return result


def max_subarray(values: list[int]) -> int:
    """Pre: nonempty integers. Post: maximum sum of a nonempty contiguous subarray."""
    if not values:
        raise ValueError("values must be nonempty")
    end = best = values[0]
    # Invariant: end is best sum ending at prior item; best is best within processed prefix.
    # Bound: remaining items.
    for value in values[1:]:
        end = max(value, end + value)
        best = max(best, end)
    return best


def demo():
    print("lower_bound([1,3,3,7],3):", lower_bound([1, 3, 3, 7], 3))
    print("sort_colors([2,0,1,2,0]):", sort_colors([2, 0, 1, 2, 0]))
    print("max_subarray([-2,3,-1,4,-5]):", max_subarray([-2, 3, -1, 4, -5]))


def self_test():
    count = 0
    # Invariant: shorter/preceding arrays match independent stdlib/brute-force results.
    # Bound: remaining arrays with length at most 5.
    for length in range(6):
        for values in product(range(3), repeat=length):
            original = list(values)
            check(sort_colors(original), sorted(original))
            check(original, list(values))
            for target in range(4):
                check(lower_bound(sorted(original), target), bisect_left(sorted(original), target))
            count += 6
    for length in range(1, 5):
        for values in product((-2, 0, 3), repeat=length):
            reference = max(sum(values[i:j]) for i in range(length) for j in range(i+1, length+1))
            check(max_subarray(list(values)), reference)
            count += 1
    return count

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()