Conditional Probability from Joint Distribution

Problem

Compute the Conditional Probability from a Joint Distribution. Given a joint probability table P(X, Y), compute the conditional distribution P(X | Y = y) or P(Y | X = x).

Solution

import numpy as np

def conditional_from_joint(joint: np.ndarray, given_axis: int,
                           given_index: int) -> np.ndarray:
    if given_axis == 1:
        joint_slice = joint[:, given_index]
    else:
        joint_slice = joint[given_index, :]

    marginal = np.sum(joint_slice)
    if marginal == 0:
        return np.zeros_like(joint_slice)
    return joint_slice / marginal

Explanation

1. The joint distribution is a 2D array where joint[i, j] = P(X=i, Y=j).
2. To compute P(X | Y=y), take the column for Y=y and divide by P(Y=y) = sum of that column.
3. To compute P(Y | X=x), take the row for X=x and divide by P(X=x) = sum of that row.
4. The result is a valid probability distribution that sums to 1 (assuming the marginal is non-zero).

Complexity

• Time: O(n) where n is the size of the relevant dimension
• Space: O(n) for the resulting conditional distribution