Hypergeometric Distribution PMF

#245 · Probability · Medium

Problem

Compute the probability mass function (PMF) of the hypergeometric distribution. Given a population of N items with K successes, compute the probability of getting exactly k successes in a sample of n drawn without replacement.

Solution

def hypergeometric_pmf(N: int, K: int, n: int, k: int) -> float:
    """
    N: population size
    K: number of success states in population
    n: number of draws
    k: number of observed successes
    """
    def log_comb(a, b):
        if b < 0 or b > a:
            return float('-inf')
        if b == 0 or b == a:
            return 0.0
        b = min(b, a - b)
        result = 0.0
        for i in range(b):
            result += log(a - i) - log(i + 1)
        return result

    from math import log, exp

    # P(X=k) = C(K,k) * C(N-K, n-k) / C(N, n)
    if k < max(0, n - (N - K)) or k > min(K, n):
        return 0.0

    log_pmf = log_comb(K, k) + log_comb(N - K, n - k) - log_comb(N, n)
    return round(exp(log_pmf), 6)

Explanation

The hypergeometric PMF is: P(X=k) = C(K,k) * C(N-K, n-k) / C(N,n).
C(K,k) counts ways to choose k successes from K success items.
C(N-K, n-k) counts ways to choose the remaining draws from failure items.
C(N,n) counts total ways to draw n items from N.
Log-space computation avoids overflow with large factorials.
Valid range check: k must be between max(0, n-(N-K)) and min(K, n).

Complexity

Time: O(min(k, K-k) + min(n-k, N-K-n+k) + min(n, N-n))
Space: O(1)

← #244 #246 →