LeetCode 🟡 2929. Distribute Candies Among Children II

🔗 🟡 2929. Distribute Candies Among Children II 1701

tags: Biweekly Contest 117 枚舉(Enumeration) 數學(Math) 組合數學(Combinatorics)

本題是 2928 的數據範圍加大版本，暴力做法可以參考該題的 解題紀錄，但在本題中使用會 TLE 。

題意

給定兩個正整數 $n$ 和 $limit$。

返回將 $n$ 個糖果分配給 $3$ 個小孩的總方式數，且每位小孩獲得的糖果不超過 $limit$。

限制條件

• $1 \leq n \leq 10^6$
• $1 \leq limit \leq 10^6$

思路

方法一：優化枚舉 + 一重迴圈

前一題 2928 中的暴力解法顯然太過簡單粗暴，讓我們更進一步優化。首先，可以先確定 $i$ 的範圍為 $[0, \min(\text{limit}, n)]$，而此時 $j + k = n - i$

• 若 $j + k = n - i > 2 \times \text{limit}$，則不管怎麼分配，$j$ 和 $k$ 兩者間必有一人獲得超過 $\text{limit}$ 個糖果，因此不合法。
• 否則，則可以考慮 $j$ 的範圍：
  • $j$ 的最大值 $mx = \min(n - i, \text{limit})$ ，即將剩餘的糖果全部給 $j$，但需滿足 $j \leq \text{limit}$。
  • $j$ 的最小值 $mn = \max(0, n - i - \text{limit})$ ，即將剩餘的糖果全部給 $k$，但需滿足 $k \leq \text{limit}$。
• 而當 $i, j$ 皆確定後，$k$ 也就確定了，因此可以得出此時 $j$ 的範圍為 $[mn, mx]$，有 $mx - mn + 1$ 種分配方式。

如此一來，我們只需枚舉 $i$，再根據 $i$ 確定 $j$ 的範圍，即可計算出答案。

複雜度分析

• 時間複雜度 $O(\min(\text{limit}, n))$。
• 空間複雜度 $O(1)$

PythonC++

class Solution:
    def distributeCandies(self, n: int, limit: int) -> int:
        ans = 0
        for i in range(min(limit, n) + 1):
            if n - i > 2 * limit:
                continue
            mx = min(n - i, limit) # j 的上限
            mn = max(0, n - i - limit) # j 的下限
            ans += mx - mn + 1
        return ans

class Solution {
public:
    int distributeCandies(int n, int limit) {
        int ans = 0, mx, mn;
        for (int i = 0; i <= min(limit, n); i++) {
            if (n - i > 2 * limit) continue;
            mx = min(n - i, limit); // j 的最大值
            mn = max(0, n - i - limit); // j 的最小值
            ans += mx - mn + 1;
        }
        return ans;
    }
};

方法二：排容原理

除了枚舉之外，也能用 排容原理 計算，即計算所有可能的組合數，再扣除不合法的情況，若有重複扣除的情況，則需要再加回來，以此類推。

而將 $n$ 個糖果分配給 $3$ 個小孩的任意方式數為 $H^3_n = \binom{n+2}{2}$，即 $C(n+2, 2)$，可以用 隔板法 思考。而我們可以用此方式繼續計算出不合法的情況，即至少一人拿到超過 $\text{limit}$ 個糖果的情況，以及至少兩人拿到超過 $\text{limit}$ 個糖果的情況。

• 初始化答案為 $H^3_n = \binom{n+2}{2}$。
• 若 $n \geq \text{limit} + 1$，則在所有方案中，存在某些方案使得至少一人拿到超過 $\text{limit}$ 個糖果，計算這種情況的方式數：
  • 計算方式為先將 $limit + 1$ 個糖果分給其中一人，則剩下 $n - (limit + 1)$ 個糖果，則有 $H^3_{n-(limit+1)} = \binom{n-(limit+1)+2}{2} = \binom{n-limit+1}{2}$ 種分配方式。
  • 由於有三個小孩，選擇 $1$ 個小孩的方法數有 $\binom{3}{1} = 3$ 種，因此有 $3 \times H^3_{n-(limit+1)} = 3 \times \binom{n-limit+1}{2}$ 種分配方式。
  • 將答案減去 $3 \times \binom{n-limit+1}{2}$。
• 若 $n \geq 2 \times (\text{limit} + 1)$，則在所有方案中，存在某些方案使得至少兩人拿到超過 $\text{limit}$ 個糖果，計算這種情況的方式數：
  • 計算方式為先將 $2 \times (limit + 1)$ 個糖果分給其中兩人，則剩下 $n - 2 \times (limit + 1)$ 個糖果，則有 $H^3_{n-2 \times (limit+1)} = \binom{n-2 \times (limit+1)+2}{2} = \binom{n-2 \times limit}{2}$ 種分配方式。
  • 由於有三個小孩，選擇 $2$ 個小孩的方法數有 $\binom{3}{2} = 3$ 種，因此有 $3 \times H^3_{n-2 \times (limit+1)} = 3 \times \binom{n-2 \times limit}{2}$ 種分配方式。
  • 將答案加上 $3 \times \binom{n-2 \times limit}{2}$。
• 若 $n \geq 3 \times (\text{limit} + 1)$，則在所有方案中，存在某些方案使得三人皆拿到超過 $\text{limit}$ 個糖果，計算這種情況的方式數：
  • 計算方式為先將 $3 \times (limit + 1)$ 個糖果分給三人，則剩下 $n - 3 \times (limit + 1)$ 個糖果，則有 $H^3_{n-3 \times (limit+1)} = \binom{n-3 \times (limit+1)+2}{2} = \binom{n-3 \times limit-1}{2}$ 種分配方式。
  • 由於有三個小孩，選擇 $3$ 個小孩的方法數有 $\binom{3}{3} = 1$ 種，因此有 $H^3_{n-3 \times (limit+1)} = \binom{n-3 \times limit-1}{2}$ 種分配方式。
  • 將答案減去 $\binom{n-3 \times limit-1}{2}$。
• 此外，當 $n > 3 \times \text{limit}$ 時，必有至少一人拿到超過 $\text{limit}$ 個糖果，不存在合法情況。對於這種情況，可以直接返回 $0$。因此上面的第 $3$ 點其實可以省略。

複雜度分析

• 時間複雜度 $O(1)$。
• 空間複雜度 $O(1)$

PythonC++

class Solution:
    def distributeCandies(self, n: int, limit: int) -> int:
        if n > 3 * limit: # 一定會有人拿到超過 limit 個糖果
            return 0
        ans = comb(n + 2, 2)
        if n >= (limit + 1): # 至少一人拿到超過 limit 個糖果
            ans -= 3 * comb(n - limit + 1, 2)
        if n >= 2 * (limit + 1): # 至少兩人拿到超過 limit 個糖果
            ans += 3 * comb(n - 2 * limit, 2)
        # if n >= 3 * (limit + 1): # 三人皆拿到超過 limit 個糖果
        #     ans -= comb(n - 3 * limit - 1, 2)
        return ans

class Solution {
public:
    int calc(int x) { // comb(x, 2)
        if (x < 0) return 0;
        return x * (x - 1) / 2;
    }
    int distributeCandies(int n, int limit) {
        if (n > 3 * limit) return 0;
        int ans = calc(n+2); // comb(n+2, 2)
        ans -= 3 * calc(n - limit + 1);
        ans += 3 * calc(n - 2 * limit);
        // ans -= calc(n - 3 * limit - 1);
        return ans; 
    }
};

---

寫在最後

Cover photo is generated by @ゴリラの素材屋さん, thanks for their work!