Compute the probability mass function of the negative binomial distribution. Given the number of required successes r, the probability of success p, and the total number of trials k, compute the probability that the r-th success occurs on the k-th trial.
def negative_binomial_pmf(k: int, r: int, p: float) -> float:
"""
k: total number of trials
r: number of successes needed
p: probability of success on each trial
Returns: P(X = k) where X is the trial on which the r-th success occurs
"""
if k < r or r <= 0 or p <= 0 or p > 1:
return 0.0
from math import log, exp
# P(X=k) = C(k-1, r-1) * p^r * (1-p)^(k-r)
# Compute in log space
def log_comb(n, m):
if m < 0 or m > n:
return float('-inf')
m = min(m, n - m)
result = 0.0
for i in range(m):
result += log(n - i) - log(i + 1)
return result
log_pmf = log_comb(k - 1, r - 1) + r * log(p) + (k - r) * log(1 - p)
return round(exp(log_pmf), 6)P(X=k) = C(k-1, r-1) * p^r * (1-p)^(k-r).