如何用Python实现n人离开房间的精准概率分布计算？

问题描述 #

房间内有n个人，每个人都有独立的离开概率。例如3人场景：A离开概率30%，B50%，C80%，结果分为0、1、2、3人离开四种，每种对应不同概率——比如3人全离开的概率是30%×50%×80%=12%；恰好2人离开的概率是30%×50%×(1-80%) + 30%×(1-50%)×80% + (1-30%)×50%×80%=79%。

需要编写Python程序，输入：

• n：整数，房间人数；
• personProb：长度为n的列表，索引k对应第k个人的离开概率。

输出：

• resultProb：长度为n+1的列表，索引k对应恰好k人离开的概率。

当前编写代码时遇到瓶颈，初始代码片段：

probList = []
for k in range(n+1): # k represents the number of people who will leave

核心问题：对于0到n的每个k，需要枚举n选k的所有组合，计算「选中的k个概率乘积 × 未选中的(1-概率)乘积」的总和。目前不知道如何高效枚举这些组合；已知k=2时可用双指针，但k>2时该怎么实现？有没有k指针方法？

解决方案 #

方法1：用itertools.combinations枚举所有组合 #

Python标准库的itertools.combinations可以直接生成n个元素中选k个的所有下标组合，无需自己实现多指针逻辑，代码简洁易读。

完整代码示例：

import itertools

def calculate_leave_probs(n, personProb):
    resultProb = [0.0] * (n + 1)
    for k in range(n + 1):
        # 生成所有选k个人的下标组合
        for combo in itertools.combinations(range(n), k):
            prob = 1.0
            for i in range(n):
                prob *= personProb[i] if i in combo else (1 - personProb[i])
            resultProb[k] += prob
    return resultProb

# 测试示例
n = 3
personProb = [0.3, 0.5, 0.8]
print(calculate_leave_probs(n, personProb))
# 输出接近 [0.07, 0.02, 0.79, 0.12]（浮点精度误差可忽略）

注意：该方法时间复杂度为O(n×C(n,k))，当n较大（比如n>20）时，组合数会急剧增长，效率很低，仅适合小n场景。

方法2：动态规划（高效解法） #

用动态规划可以避免枚举所有组合，时间复杂度降至O(n²)，适合大n的场景。

思路：

• 定义dp[i][j]表示前i个人中恰好j人离开的概率；
• 初始状态：dp[0][0] = 1.0（0个人时，0人离开的概率为1）；
• 状态转移：对于第i个人（从1到n），有两种选择：
  1. 他离开：dp[i][j] += dp[i-1][j-1] * personProb[i-1]（注意personProb是0基下标）；
  2. 他不离开：dp[i][j] += dp[i-1][j] * (1 - personProb[i-1])。

代码示例：

def calculate_leave_probs_dp(n, personProb):
    # 初始化dp数组，dp[i][j]表示前i个人恰好j人离开的概率
    dp = [[0.0]*(n+1) for _ in range(n+1)]
    dp[0][0] = 1.0
    
    for i in range(1, n+1):
        prob = personProb[i-1]
        # 前i个人最多i人离开，遍历j的可能取值
        for j in range(0, i+1):
            # 情况1：第i个人不离开
            dp[i][j] += dp[i-1][j] * (1 - prob)
            # 情况2：第i个人离开（需保证j>=1）
            if j >= 1:
                dp[i][j] += dp[i-1][j-1] * prob
    # 最终结果为前n个人的所有可能情况
    return dp[n]

# 测试示例
n = 3
personProb = [0.3, 0.5, 0.8]
print(calculate_leave_probs_dp(n, personProb))
# 输出 [0.07, 0.02, 0.79, 0.12]（浮点精度下的准确值）

关于多指针的问题 #

如果不想依赖itertools，自己实现多指针枚举组合确实可行，但逻辑复杂且容易出错——比如k=3时需要三个指针控制选中元素的下标，还要保证下标递增以避免重复组合，代码量极大且维护成本高，完全没必要。优先使用itertools或动态规划方案即可。

内容的提问来源于stack exchange，提问作者Bao