Python-120：摇骰子的胜利概率

问题描述

小U和小S正在玩一个有趣的骰子游戏。每个骰子都有固定数量的面数k，每一面的点数分别是1到k。小U拥有n个骰子，每个骰子i的面数是 a_i，摇到每一面的概率均为 1/a_i。小S则有m个骰子，每个骰子j的面数是 b_j，摇到每一面的概率均为 1/b_j。

两人分别同时摇各自的骰子，并将摇出的点数相加，得分较高的一方获胜，得分相同则为平局。游戏只进行一次，没有重赛。现在小U想知道他获胜的概率是多少。你能帮他计算吗？（答案保留三位小数）

输入：

• n：小U的骰子数量
• m：小S的骰子数量
• arrayN：小U的每个骰子面数
• arrayM：小S的每个骰子面数

约束条件：

• n, m, k为整数，数据范围 1 ≤ n, m, k ≤ 25

代码

def solution(n, m, arrayN, arrayM):

    # Edit your code here

    assert n == len(arrayN)

    assert m == len(arrayM)

    def gen_table(a):

        s = sum(a)

        f = [0 for _ in range(s + 1)]

        f[0] = 1

        for x in a:

            g = [0 for _ in range(s + 1)]

            for i, v in enumerate(f):

                for j in range(1, x + 1):

                    if i + j <= s:

                        g[i + j] += v / x

            f = g

        return f

    p1 = gen_table(arrayN)

    p2 = gen_table(arrayM)

    ans = 0

    for i in range(len(p1)):

        for j in range(len(p2)):

            if j < i:

                ans += p1[i] * p2[j]

    return round(ans, 3)

 

if __name__ == "__main__":

    # Add your test cases here

    print(solution(1, 3, [8], [2, 3, 4]) == 0.255)

    print(solution(2, 2, [3, 4], [3, 3]) == 0.5)

    print(solution(3, 1, [2, 2, 2], [4]) == 0.844)