binary search memo

bs.py

from typing import List
import bisect


def binary_search(a: List[int], num: int):
    min_idx = 0
    max_idx = len(a) - 1
    while min_idx <= max_idx:
        mid_idx = (min_idx + max_idx) // 2
        curr = a[mid_idx]
        if curr == num:
            return mid_idx
        
        if curr > num:
            max_idx = mid_idx - 1
        else:
            min_idx = mid_idx + 1

    return -1

def binary_search_left(a: List[int], num: int):
    min_idx = 0
    max_idx = len(a) - 1
    while min_idx <= max_idx:
        mid_idx = (min_idx + max_idx) // 2
        curr = a[mid_idx]
        if curr == num:
            return mid_idx
        
        if curr > num:
            max_idx = mid_idx - 1
        else:
            min_idx = mid_idx + 1

    return max_idx


def binary_search_right(a: List[int], num: int):
    min_idx = 0
    max_idx = len(a) - 1
    while min_idx <= max_idx:
        mid_idx = (min_idx + max_idx) // 2
        curr = a[mid_idx]
        if curr == num:
            return mid_idx
        
        if curr > num:
            max_idx = mid_idx - 1
        else:
            min_idx = mid_idx + 1

    return min_idx


def main():
    # len(a) = 10
    a = [1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

    assert binary_search(a, 1) == 0
    assert binary_search(a, 6) == 2
    assert binary_search(a, 55) == 9

    assert binary_search(a, 0) == -1
    assert binary_search(a, 7) == -1
    assert binary_search(a, 56) == -1

    assert binary_search_left(a, 1) == 0
    assert binary_search_left(a, 6) == 2
    assert binary_search_left(a, 55) == 9

    assert binary_search_left(a, 0) == -1
    assert binary_search_left(a, 7) == 2
    assert binary_search_left(a, 56) == 9

    assert binary_search_right(a, 1) == 0
    assert binary_search_right(a, 6) == 2
    assert binary_search_right(a, 55) == 9

    # 0
    assert binary_search_right(a, 0) == bisect.bisect(a, 0)
    # 3
    assert binary_search_right(a, 7) == bisect.bisect(a, 7)
    # 10
    assert binary_search_right(a, 56) == bisect.bisect(a, 56)

if __name__ == '__main__':
    main()