impl/python.py

from math import gcd
from typing import List, Sequence

from .stats_mean_median_types import CentralTendency, Fraction


MAX_SAFE = 9007199254740991


def _fraction(numerator: int, denominator: int) -> Fraction:
    g = gcd(numerator, denominator) or 1
    return Fraction(numerator=numerator // g, denominator=denominator // g)


def mean_median_mode(values: Sequence[int]) -> CentralTendency:
    """Mean, median and mode of a list of integers.

    Mean and median are exact fractions; their float forms come from a single
    correctly rounded division, so every language returns the same double.
    """
    if isinstance(values, (str, bytes)) or not isinstance(values, (list, tuple)):
        raise TypeError("values must be a list of integers")
    if len(values) == 0:
        raise ValueError("values must not be empty")
    for v in values:
        # bool is an int in Python; every other language rejects it, so do too.
        if isinstance(v, bool) or not isinstance(v, int):
            raise TypeError("values must be integers, received %r" % (v,))
        if abs(v) > MAX_SAFE:
            raise ValueError("values must be safe integers (magnitude at most 9007199254740991), received %d" % v)

    total = sum(values)
    if abs(total) > MAX_SAFE:
        raise ValueError("sum exceeds the safe integer range")
    count = len(values)

    ordered = sorted(values)
    if count % 2 == 1:
        median = Fraction(numerator=ordered[(count - 1) // 2], denominator=1)
    else:
        pair = ordered[count // 2 - 1] + ordered[count // 2]
        if pair % 2 == 0:
            median = Fraction(numerator=pair // 2, denominator=1)
        else:
            if abs(pair) > MAX_SAFE:
                raise ValueError("median exceeds the safe integer range")
            median = Fraction(numerator=pair, denominator=2)

    # Runs of equal values in the sorted copy give frequencies in ascending
    # value order, so the modes come out sorted without a second sort.
    modes: List[int] = []
    mode_frequency = 0
    i = 0
    while i < count:
        j = i
        while j < count and ordered[j] == ordered[i]:
            j += 1
        run = j - i
        if run > mode_frequency:
            mode_frequency = run
            modes = [ordered[i]]
        elif run == mode_frequency:
            modes.append(ordered[i])
        i = j

    # float(a) / float(b) with both exact (|a| <= 2^53) is one correctly
    # rounded IEEE division, the same double TypeScript and Rust produce.
    return CentralTendency(
        count=count,
        sum=total,
        mean=_fraction(total, count),
        mean_value=float(total) / float(count) + 0.0,
        median=median,
        median_value=float(median.numerator) / float(median.denominator) + 0.0,
        modes=modes,
        mode_frequency=mode_frequency,
    )