Made for tutoring lesson

mean_std.py

"""Calculate the mean, variance, and standard deviation of a list of numbers."""

__author__ = "Jesse Wei <jessew13@email.unc.edu>"

EPSILON = 0.01

ANSWERS = {
        'A': [1, 1, 1, 10, 10, 10],
        'B': [1, 2, 3, 4, 5, 6],
        'C': [2, 6, 6, 10, 10, 12],
        'D': [5, 5, 5, 5, 5, 5, 5],
        'E': [5, 6, 7, 8, 9, 10]
        }

E = {
        5: 0.3,
        6: 0.1,
        7: 0.1,
        8: 0.1,
        9: 0.1,
        10: 0.3
        }


def main() -> None:
    for answer in ANSWERS:
        print(f"Var({answer}) = {var(ANSWERS[answer])}")
    # print(f"Var(E) = {var_pmf(E)}")
    

def mean(xs: list[int]) -> float:
    """Return the mean of a list."""
    return sum(xs) / len(xs)


def var(xs: list[int]) -> float:
    """Return the variance of a list."""
    variance = 0
    mu = mean(xs)
    for elt in xs:
        variance += (elt - mu) ** 2 / len(xs)
    return variance


def mean_pmf(xs) -> float:
    """Return the mean, or expected value, of a probability mass function."""
    assert abs(sum(xs.values()) - 1) < EPSILON
    mean = 0
    for elt in xs:
        mean += elt * xs[elt]
    return mean


def var_pmf(xs: dict[int, float]) -> float:
    """Return the variance of a probability mass function."""
    variance = 0
    mu = mean_pmf(xs)
    for elt in xs:
        variance += (elt - mu) ** 2 * xs[elt]
    return variance


if __name__ == "__main__":
    main()