LeetCode每日一题2025-06-01

点击获取AI摘要

2929. 给小朋友们分糖果 II M

给你两个正整数 n 和 limit 。

请你将 n 颗糖果分给 3 位小朋友，确保没有任何小朋友得到超过 limit 颗糖果，请你返回满足此条件下的 总方案数 。

示例 1：

输入：n = 5, limit = 2
输出：3
解释：总共有 3 种方法分配 5 颗糖果，且每位小朋友的糖果数不超过 2 ：(1, 2, 2) ，(2, 1, 2) 和 (2, 2, 1) 。

示例 2：

输入：n = 3, limit = 3
输出：10
解释：总共有 10 种方法分配 3 颗糖果，且每位小朋友的糖果数不超过 3 ：(0, 0, 3) ，(0, 1, 2) ，(0, 2, 1) ，(0, 3, 0) ，(1, 0, 2) ，(1, 1, 1) ，(1, 2, 0) ，(2, 0, 1) ，(2, 1, 0) 和 (3, 0, 0) 。

提示：

• 1 <= n <= 10⁶
• 1 <= limit <= 10⁶

问题分析

我们需要计算满足以下条件的非负整数解 $(x, y, z)$ 的个数：

1. $x + y + z = n$；
2. $0 \le x, y, z \le \mathrm{limit}$。

如果不考虑上界限制，则“将 $n$ 个相同的糖果分给 3 位小朋友”的非负整数解数量为星与条公式：

$$\binom{n + 3 - 1}{3 - 1} \;=\; \binom{n + 2}{2}\;=\;\frac{(n+2)(n+1)}{2}.$$

但由于每位小朋友最多只能拿到 $\mathrm{limit}$ 颗，需要排除掉至少有一个孩子拿到超过 $\mathrm{limit}$ 颗的情况。可通过包容–排斥原理计算满足上界约束的解的个数。

设事件：

• $A_x$：满足 $x > \mathrm{limit}$ 的解；
• $A_y$：满足 $y > \mathrm{limit}$ 的解；
• $A_z$：满足 $z > \mathrm{limit}$ 的解。

目标是计算：

$$\Bigl|\{(x,y,z)\mid x+y+z=n,\;0\le x,y,z\le \mathrm{limit}\}\Bigr| \;=\; \binom{n+2}{2} \;-\; \bigl|A_x \cup A_y \cup A_z\bigr|.$$

其中

$$\bigl|A_x \cup A_y \cup A_z\bigr| \;=\; |A_x| + |A_y| + |A_z| \;-\;\bigl(|A_x\cap A_y| + |A_x\cap A_z| + |A_y\cap A_z|\bigr) \;+\;|A_x\cap A_y\cap A_z|.$$

由于对称性，

$$|A_x|=|A_y|=|A_z|,\quad |A_x\cap A_y|=|A_x\cap A_z|=|A_y\cap A_z|.$$

• 当 $x > \mathrm{limit}$ 时，令 $x' = x - (\mathrm{limit} + 1)$，则 $x' \ge 0$，原方程变为

  $$x' + y + z \;=\; n - (\mathrm{limit} + 1).$$

  此时 $x',y,z \ge 0$，无上界，解的个数为

  $$\binom{\,(n - (\mathrm{limit}+1)) + 2}{2},$$

  仅当 $n - (\mathrm{limit}+1) \ge 0$ 时才非零，否则视作 0。

• 当同时 $x>\mathrm{limit}$ 且 $y>\mathrm{limit}$ 时，令

  $$x' = x - (\mathrm{limit}+1),\quad y' = y - (\mathrm{limit}+1),$$

  则

  $$x' + y' + z \;=\; n - 2(\mathrm{limit}+1),$$

  解的个数为

  $$\binom{\,(n - 2(\mathrm{limit}+1)) + 2}{2},$$

  条件是 $n - 2(\mathrm{limit}+1) \ge 0$。

• 当 $x,y,z$ 都超过 $\mathrm{limit}$ 时，令

  $$x' = x - (\mathrm{limit}+1),\; y' = y - (\mathrm{limit}+1),\; z' = z - (\mathrm{limit}+1),$$

  则

  $$x' + y' + z' \;=\; n - 3(\mathrm{limit}+1),$$

  解的个数为

  $$\binom{\,(n - 3(\mathrm{limit}+1)) + 2}{2},$$

  条件是 $n - 3(\mathrm{limit}+1) \ge 0$。

由包容–排斥得到：

$$\bigl|A_x \cup A_y \cup A_z\bigr| = 3\,\binom{\,n - (\mathrm{limit}+1) + 2}{2} - 3\,\binom{\,n - 2(\mathrm{limit}+1) + 2}{2} + \binom{\,n - 3(\mathrm{limit}+1) + 2}{2},$$

（不足条件时对应组合数取 0）。因此，合法分配数：

$$\boxed{ \mathrm{Ans} \;=\; \binom{n+2}{2} \;-\;3\,\binom{\,n - (\mathrm{limit}+1) + 2}{2} \;+\;3\,\binom{\,n - 2(\mathrm{limit}+1) + 2}{2} \;-\;\binom{\,n - 3(\mathrm{limit}+1) + 2}{2} }.$$

算法思路

1. 定义辅助函数 $\text{comb2}(m)$，用于计算

   $$\binom{m + 2}{2} \;=\; \frac{(m+2)(m+1)}{2},$$

   若 $m < 0$，则返回 0，以便处理 $n - k(\mathrm{limit}+1) < 0$ 时组合数为 0 的情形。
2. 计算无约束总数：

   $$\text{total} \;=\; \text{comb2}(n).$$

3. 令 $\text{step} = \mathrm{limit} + 1$，计算包容–排斥各项：
   • $\text{term1} = 3\,\text{comb2}\bigl(n - \text{step}\bigr)$；
   • $\text{term2} = 3\,\text{comb2}\bigl(n - 2\,\text{step}\bigr)$；
   • $\text{term3} = \text{comb2}\bigl(n - 3\,\text{step}\bigr)$。
4. 返回 $\text{total} - \text{term1} + \text{term2} - \text{term3}$ 即为最终结果。

时间复杂度

• 该方法只进行常数次的组合数计算和加减乘除运算，整个算法的时间复杂度为 $O(1)$，满足 $n, \mathrm{limit}$ 最大到 $10^6$ 的情况。
• 空间复杂度为 $O(1)$，只使用常数个整型变量。

代码分解

1. 辅助函数 comb2(m)

   • 功能：若 $m \ge 0$，返回 $\binom{m+2}{2} = \frac{(m+2)(m+1)}{2}$；否则返回 0。
   • 作用：统一处理“$n - k(\mathrm{limit}+1) < 0$ 时的组合数为 0”的情形。

2. 主函数 distributeCandies(n, limit)

   • 计算无上界总方案数 total = comb2(n)。
   • 设 step = limit + 1。
   • 计算：
     • term1 = 3 * comb2(n - step) 对应排除掉任一孩子超过上限的情况；
     • term2 = 3 * comb2(n - 2*step) 对应排除掉任意两位孩子都超过上限的情况；
     • term3 = comb2(n - 3*step) 对应排除掉三位孩子都超过上限的情况。
   • 最终返回 total - term1 + term2 - term3。

代码实现

class Solution:
    def distributeCandies(self, n: int, limit: int) -> int:
        # 辅助函数：计算 C(m+2, 2)，m < 0 时返回 0
        def comb2(m: int) -> int:
            if m < 0:
                return 0
            # C(m+2, 2) = (m+2)*(m+1)//2
            return (m + 2) * (m + 1) // 2

        # 无上界时的总解数：C(n+2, 2)
        total = comb2(n)

        # 每个孩子超过上界时，需要从 n 中减去 (limit+1)
        step = limit + 1

        # 包容-排斥各项
        term1 = 3 * comb2(n - step)       # 单个孩子超过上限
        term2 = 3 * comb2(n - 2 * step)   # 任意两位孩子都超过上限
        term3 =     comb2(n - 3 * step)   # 三位孩子都超过上限

        # 最终结果
        return total - term1 + term2 - term3