median_in_a_stream.py

import heapq


def signum(a: int, b: int) -> int:
    """
    Compare two integers.

    Returns:
        1 if a > b
       -1 if a < b
        0 if a == b
    """
    if a > b:
        return 1
    if a < b:
        return -1
    return 0


def call_median(
    element: int,
    max_heap: list[int],
    min_heap: list[int],
    median: int,
) -> int:
    """
    Insert an element into heaps and update the median.
    """
    case = signum(len(max_heap), len(min_heap))

    if case == 0:
        if element > median:
            heapq.heappush(min_heap, element)
            median = min_heap[0]
        else:
            heapq.heappush(max_heap, -element)
            median = -max_heap[0]

    elif case == 1:
        if element > median:
            heapq.heappush(min_heap, element)
        else:
            heapq.heappush(min_heap, -heapq.heappop(max_heap))
            heapq.heappush(max_heap, -element)
        median = (-max_heap[0] + min_heap[0]) // 2

    else:
        if element > median:
            heapq.heappush(max_heap, -heapq.heappop(min_heap))
            heapq.heappush(min_heap, element)
        else:
            heapq.heappush(max_heap, -element)
        median = (-max_heap[0] + min_heap[0]) // 2

    return median


def median_in_a_stream(numbers: list[int]) -> list[int]:
    """
    Find the median after each insertion in a stream of integers.

    Uses two heaps and follows the classic running median logic.

    Reference:
    https://en.wikipedia.org/wiki/Median#Running_median

    Args:
        numbers: List of integers

    Returns:
        List of medians after each insertion

    Raises:
        ValueError: If the input list is empty

    >>> median_in_a_stream([20, 14, 13, 16, 17])
    [20, 17, 14, 15, 16]
    >>> median_in_a_stream([5, 15, 1, 3])
    [5, 10, 5, 4]
    >>> median_in_a_stream([])
    Traceback (most recent call last):
    ...
    ValueError: Input list must not be empty
    """
    if not numbers:
        raise ValueError("Input list must not be empty")

    max_heap: list[int] = []
    min_heap: list[int] = []
    median = 0
    result: list[int] = []

    for element in numbers:
        median = call_median(element, max_heap, min_heap, median)
        result.append(median)

    return result