多颜色取球条件概率计算的代码优化方法问询

问题背景 #

盒子内装有12种颜色共100个球，规则如下：

• 最多有3次抽取机会，目标是抽到特定颜色的球
• 若抽到非目标颜色的球，立即将该颜色的所有球从盒中移除
• 手动枚举所有颜色组合计算成功概率过于冗余，需通用高效的实现方案

原手动实现代码 #

def calx(prob1, prob2, prob3, prob4):
    first_selection = prob1
    second_selection = prob2*(prob1/(1-prob2))+prob3*(prob1/(1-prob3))+prob4*(prob1/(1-prob4))
    third_selection = prob2*(prob3/(1-prob2))*(prob1/(1-prob2-prob3))
                      +prob3*(prob2/(1-prob3))*(prob1/(1-prob2-prob3))
                      +prob3*(prob4/(1-prob3))*(prob1/(1-prob3-prob4))
                      +prob4*(prob3/(1-prob4))*(prob1/(1-prob3-prob4))
                      +prob2*(prob4/(1-prob2))*(prob1/(1-prob2-prob4))
                      +prob4*(prob2/(1-prob4))*(prob1/(1-prob2-prob4))

    return first_selection+second_selection+third_selection

calx(.5,.4,.05,.05)

通用数学公式（翻译自英文原版） #

设目标颜色的抽取概率为 ( p_t )，其余11种颜色的抽取概率分别为 ( p_1, p_2, ..., p_{11} )，则成功抽到目标颜色的总概率为：
[
P = p_t + \sum_{i=1}^{11} \left( \frac{p_t}{1-p_i} \cdot p_i \right) + \sum_{1 \leq i < j \leq 11} \left( \frac{p_t}{1-p_i-p_j} \cdot \frac{p_j}{1-p_i} \cdot p_i + \frac{p_t}{1-p_i-p_j} \cdot \frac{p_i}{1-p_j} \cdot p_j \right)
]
公式拆解：

• 第一项：第一次抽取即命中目标的概率
• 第二项：第一次抽到非目标颜色( i )后，第二次命中目标的概率之和
• 第三项：前两次抽到不同非目标颜色( i,j )（两种顺序）后，第三次命中目标的概率之和

高效通用实现代码 #

利用itertools.combinations自动生成所有非目标颜色的组合，无需手动枚举，支持任意数量的非目标颜色：

import itertools

def calculate_success_prob(target_prob, other_probs):
    # 第一次命中目标的概率
    first = target_prob
    
    # 第二次命中的概率：第一次抽中非目标，第二次命中
    second = 0.0
    for p in other_probs:
        denom = 1 - p
        if denom == 0:
            continue  # 避免除零，实际场景中不会出现该情况
        second += p * (target_prob / denom)
    
    # 第三次命中的概率：前两次抽中不同非目标，第三次命中
    third = 0.0
    # 生成所有两两不同的非目标颜色组合
    for p_i, p_j in itertools.combinations(other_probs, 2):
        denom_i = 1 - p_i
        denom_j = 1 - p_j
        denom_both = 1 - p_i - p_j
        if denom_i == 0 or denom_j == 0 or denom_both == 0:
            continue
        # 顺序1：先抽中p_i，再抽中p_j，最后命中目标
        term1 = p_i * (p_j / denom_i) * (target_prob / denom_both)
        # 顺序2：先抽中p_j，再抽中p_i，最后命中目标
        term2 = p_j * (p_i / denom_j) * (target_prob / denom_both)
        third += term1 + term2
    
    return first + second + third

# 示例调用：对应原代码参数，目标概率0.5，其他颜色概率[0.4, 0.05, 0.05]
print(calculate_success_prob(0.5, [0.4, 0.05, 0.05]))

实现优势 #

• 自动生成所有组合，无需手动编写每个颜色的排列逻辑
• 支持任意数量的非目标颜色（如题目中的11种），扩展性强
• 加入除零判断，避免运行时异常

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