Cálculo da variância em Python

Este post segue uma seqüência iniciada no blog Kodumaro, comparando a implementação de alguns cálculos de estatística usando linguagens como Lisp, C e depois seguido de Ruby no blog Programando sem cafeína, e Smalltalk.

Como estou vendo um pouco de estatística na faculdade agora, me animei para tentar uma implementação em Python também.

import math

def arithmetic_mean(dataset):
    """
    Simple arithmetic mean (average)
    
    Example:
    >>> arithmetic_mean([10,100,1000])
    370
    """
    return sum(dataset) / len(dataset)


def root_mean_square(dataset):
    """
    RMS (quadratic mean)

    Example:
    >>> root_mean_square([10,100,1000])
    580.25856305616037
    """
    return math.sqrt(sum(x**2 for x in dataset) / len(dataset))


def var_num(dataset, mean_function=arithmetic_mean):
    """
    Divisor for variance calculation
    """
    mean = mean_function(dataset)
    return sum((x - mean)**2 for x in dataset)


def populational_variance(dataset):
    """
    Population variance
    
    Example:
    >>> populational_variance([10,100,1000])
    199800
    """
    return var_num(dataset) / len(dataset)


def sample_variance(dataset):
    """
    Sample variance
    
    Example:
    >>> sample_variance([10,100,1000])
    299700
    """
    return var_num(dataset) / (len(dataset) - 1)

Preferi implementar como funções simples pois assim é fácil de aplicar sobre qualquer conjunto de dados iterável como listas, tuplas, dicionários ou sets.

Na função de numerador da variância, um dos parâmetros é a função de média que se deseja utilizar para o cálculo, com o padrão sendo a média aritmética (pois como nota o blog Kodumaro, o uso da função depende de quem pediu o cálculo ;).

Vale notar também o uso de generator expressions para os somatórios, que substituem as mais tradicionais list comprehensions nesse caso por questão de performance, principalmente para grandes conjuntos de dados.

Apliquei ainda doctests simples para validar os resultados das funções.