← back

Compute Covariance from Joint Probability Mass Function

#243 · Probability · Medium

⊣ Solve on deep-ml.com

Problem

Compute the covariance of two discrete random variables given their joint probability mass function (PMF).

Solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def covariance_from_joint_pmf(
    x_values: list[float],
    y_values: list[float],
    joint_pmf: list[list[float]],
) -> float:
    """
    x_values: possible values of X
    y_values: possible values of Y
    joint_pmf: [len(x_values), len(y_values)] where joint_pmf[i][j] = P(X=x_i, Y=y_j)
    """
    # E[X]
    e_x = 0.0
    for i, x in enumerate(x_values):
        marginal_x = sum(joint_pmf[i][j] for j in range(len(y_values)))
        e_x += x * marginal_x

    # E[Y]
    e_y = 0.0
    for j, y in enumerate(y_values):
        marginal_y = sum(joint_pmf[i][j] for i in range(len(x_values)))
        e_y += y * marginal_y

    # E[XY]
    e_xy = 0.0
    for i, x in enumerate(x_values):
        for j, y in enumerate(y_values):
            e_xy += x * y * joint_pmf[i][j]

    # Cov(X, Y) = E[XY] - E[X]E[Y]
    cov = e_xy - e_x * e_y
    return round(cov, 6)

Explanation

  1. Compute marginal distributions by summing the joint PMF along each axis.
  2. Calculate E[X] and E[Y] from the marginals.
  3. Calculate E[XY] by summing x y P(X=x, Y=y) over all pairs.
  4. Apply the covariance formula: Cov(X, Y) = E[XY] - E[X] * E[Y].

Complexity

  • Time: O(m * n) where m and n are the sizes of x_values and y_values
  • Space: O(1) beyond the input
← #242#244 →