Compute Total Probability using Law of Total Probability

Problem

Compute the total probability of an event using the Law of Total Probability. Given conditional probabilities P(A|B_i) and partition probabilities P(B_i), compute P(A).

Solution

def total_probability(
    p_b: list[float],
    p_a_given_b: list[float],
) -> float:
    """
    p_b: [P(B_1), P(B_2), ..., P(B_n)] - partition probabilities (should sum to 1)
    p_a_given_b: [P(A|B_1), P(A|B_2), ..., P(A|B_n)]
    Returns: P(A)
    """
    p_a = sum(pb * pab for pb, pab in zip(p_b, p_a_given_b))
    return round(p_a, 6)

Explanation

1. The Law of Total Probability states: P(A) = sum_i P(A|B_i) * P(B_i) where {B_i} is a partition of the sample space.
2. Each term P(A|B_i) * P(B_i) represents the probability of A occurring through the "pathway" B_i.
3. The partition must be mutually exclusive and exhaustive: P(B_1) + P(B_2) + ... + P(B_n) = 1.
4. This is fundamental for computing marginal probabilities and forms the basis of Bayesian reasoning.

Complexity

• Time: O(n) where n is the number of partition elements
• Space: O(1)