请教CodeWars骰子点数组合问题的Python代码逻辑

在CodeWars上碰到一道编程题，自己没能独立解决，看了已通过的解决方案但完全看不懂，ChatGPT的解释前后矛盾，特此求助理解以下代码的逻辑：

def outcome(n, s, k):
    from math import comb
    if not k: return 1
    if not n or k < n: return 0
    return sum((-1) ** i * comb(n, i) * comb(k - s * i - 1, k - s * i - n) for i in range((k - n) // s + 1))

题目描述 #

有n个s面骰子，每个骰子的点数为1到s，求点数和为指定数值k的结果组合数。
例如，投掷4个6面骰子，点数和为5的组合有4种：

(1, 1, 1, 2)
(1, 1, 2, 1)
(1, 2, 1, 1)
(2, 1, 1, 1)

测试用例范围 #

0 <= n <= 10
1 <= s <= 20
0 <= k <= n * s

注意事项 #

• 无论骰子数量多少，k=0时总有1种情况；
• 没有骰子时，无法得到任何正数值k。

我最初想用递归解决，但未能成功，论坛查找也无果。

代码逻辑拆解 #

1. 导入工具函数 #

from math import comb：comb(a, b)用于计算组合数，即从a个元素中选b个的无顺序选法数量，公式为C(a,b) = a!/(b!*(a-b)!)。

2. 边界条件处理 #

• if not k: return 1：对应题目注意事项，k=0时直接返回1种情况；
• if not n or k < n: return 0：
  • not n表示没有骰子，无法得到正k，返回0；
  • k < n：每个骰子最小点数是1，n个骰子的最小和是n，k比n小不可能达到，返回0。

3. 核心计算：容斥原理的应用 #

这部分是代码的核心，用容斥原理计算满足条件的组合数，步骤如下：

问题转化

原问题是求x₁+x₂+...+xₙ = k的有序解个数，其中每个xᵢ满足1 ≤ xᵢ ≤ s（每个骰子的点数范围）。

先做变量替换：令yᵢ = xᵢ - 1，则yᵢ ≥ 0，且yᵢ ≤ s-1，原等式变为：
y₁+y₂+...+yₙ = k - n

现在问题转化为：求上述等式的非负整数解个数，且每个yᵢ ≤ s-1。

无约束解的数量

如果不考虑yᵢ ≤ s-1的限制，非负整数解的个数是组合数C((k-n)+n-1, n-1)，也就是C(k-1, n-1)。利用组合数性质C(a,b)=C(a,a-b)，可以写成C(k-1, k-n)，这对应代码中i=0时的项（(-1)^0=1，comb(n,0)=1，此时comb(k-s*0-1, k-s*0-n)=comb(k-1, k-n)）。

容斥原理修正

我们需要减去那些存在至少一个yᵢ ≥ s的非法解，再加回减去过多的部分，以此类推：

• 选i个骰子，让它们的yᵢ ≥ s（对应原骰子点数xᵢ ≥ s+1，超出范围）；
• 对这i个变量，令zᵢ = yᵢ - s，则zᵢ ≥0，此时等式变为z₁+...+zₙ = (k-n) - i*s；
• 此时解的个数为C( (k-n -i*s) +n-1, n-1 ) = C(k - s*i -1, n-1)，同样用组合数性质转化为C(k - s*i -1, k - s*i -n)，和代码中的对应项一致；
• 乘以(-1)^i是容斥的符号交替规则，乘以comb(n,i)是选i个骰子的组合数。

i的取值范围

要保证(k-n) - i*s ≥0（否则等式无意义，解数为0），即i ≤ (k-n)//s，所以i的取值范围是0到(k-n)//s，对应代码中的range((k-n)//s +1)（因为range是左闭右开，所以加1）。

把所有项求和，就得到了满足条件的组合数。

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