🔗 🟡 2597. The Number of Beautiful Subsets 2023

tags: `Weekly Contest 337` `回溯(Backtracking)` `剪枝(Pruning)` `子集型回溯` `動態規劃(Dynamic Programming)` `雜湊表(Hash Table)` `數學(Math)` `乘法原理`

題意

給定一個正整數陣列 $nums$ 和一個正整數 $k$ 。

如果 $nums$ 的子集中，任意兩個整數的絕對差值不等於 $k$ ，則該子集是一個 美麗(beautiful) 子集。

返回陣列 $nums$ 中非空且美麗的子集的數量。

限制

$1 \leq nums.length \leq 20$
$1 \leq nums[i], k \leq 1000$

思路

解法一：回溯(Backtracking) + 剪枝(Pruning)

要構造出所有可能的子集，可以使用回溯法，類似 78. Subsets，但是需要加入條件判斷。

對於遍歷到的每個元素 $x = nums[i]$ ，都有兩種可能：在當前子集中或不在當前子集中。但若 $nums[i] \pm k$ 在子集中，則不能選擇 $x$ 。為此我們需要維護當前在子集中的元素數量，可以使用雜湊表或陣列 $cnt$ 來記錄。

若不選擇 $x$ ，則直接遞迴到下一個元素，即 $dfs(i+1)$ 。
若選擇 $x$ ，則需要檢查 $cnt[x-k]$ 和 $cnt[x+k]$ 是否為 $0$ ，若不為 $0$ 則不能選擇 $x$ 。若滿足條件，需要將 $x$ 加入到子集中，即累加 $cnt[x]$ ，再遞迴到下一個元素，遞迴結束後，需要將 $x$ 從子集中移除，即減去 $cnt[x]$ 。

根據題意，會使用到的索引範圍在 $[-k, mx+k]$ 之內，其中 $mx$ 為陣列中的最大值。或是根據限制條件中 $1 \leq nums[i], k \leq 1000$ ，因此值域應為 $[-999, 2000]$ ，可以使用一個大小為 $3000$ 的陣列 $cnt$ 來記錄。但需要注意負數的情況：

在 Python 中，若索引為負數，則會從陣列的最後一個元素開始，因此只要確保後 $k$ 個索引不會被使用即可。
在 C++ 中，若索引為負數，則會造成陣列越界，因此需要將索引加上 $k$ 來避免負數索引。

而子集型回溯有兩種寫法：

選或不選當前元素，當 $i$ 等於陣列的長度時，表示已經遍歷完所有元素，此時累加答案。
枚舉下一個選的元素 $j$ ，此時每個遞迴都恰好是一個子集，因此直接在遞迴開始時累加答案即可。

在選擇時直接根據 $x$ 的值來判斷是否選擇，可以減少一些不必要的遞迴，即為 剪枝(Pruning)。

複雜度分析

時間複雜度： $\mathcal{O}(2^n)$ ，每個元素都有 $2$ 種可能，共有 $n$ 個元素，因此共有 $2^n$ 種可能。
空間複雜度： $\mathcal{O}(n)$ ，遞迴的深度為 $O(n)$ ，若使用雜湊表， $cnt$ 的空間複雜度為 $O(n)$ ；若使用陣列， $cnt$ 的空間複雜度為 $O(n+2k)$ ，本題中仍然可視為 $O(n)$ 。

寫法一：選或不選當前元素

class Solution:
    def beautifulSubsets(self, nums: List[int], k: int) -> int:
        n = len(nums)
        ans = 0 # 去除空集合
        cnt = Counter()
        # cnt = [0] * (max(nums) + 2 * k + 1)
        def dfs(i: int) -> None: # 第一種：選或不選
            nonlocal ans
            if i == n:
                ans += 1
                return
            dfs(i + 1) # 不選
            x = nums[i]
            if cnt[x-k] == 0 and cnt[x+k] == 0: # 滿足條件
                cnt[x] += 1
                dfs(i + 1) # 選
                cnt[x] -= 1
        dfs(0)
        return ans - 1 # 去除空集合

class Solution {
public:
    int beautifulSubsets(vector<int>& nums, int k) {
        int n = nums.size();
        int ans = 0, mx = *max_element(nums.begin(), nums.end());
        vector<int> cnt(mx + 2 * k + 1, 0);
        function<void(int)> dfs = [&](int i) { // 第一種：選或不選
            if (i == n) {
                ans += 1;
                return;
            }
            dfs(i + 1); // 不選
            int x = nums[i] + k; // 避免索引為負數
            if (cnt[x - k] == 0 && cnt[x + k] == 0) { // 滿足條件
                cnt[x] += 1;
                dfs(i + 1); // 選
                cnt[x] -= 1;
            }
        };
        dfs(0);
        return ans - 1; // 減去空集合
    }
};

寫法二：枚舉下一個選擇

class Solution:
    def beautifulSubsets(self, nums: List[int], k: int) -> int:
        n = len(nums)
        ans = 0
        cnt = Counter()
        # cnt = [0] * (max(nums) + 2 * k + 1)
        def dfs(i: int) -> None: # 第二種：枚舉下一個選哪個
            nonlocal ans
            ans += 1
            for j in range(i, n): # 枚舉下一個選的是哪個
                x = nums[j]
                if cnt[x-k] == 0 and cnt[x+k] == 0: # 滿足條件
                    cnt[x] += 1
                    dfs(j + 1)
                    cnt[x] -= 1
        dfs(0)
        return ans - 1 # 去除空集合

class Solution {
public:
    int beautifulSubsets(vector<int>& nums, int k) {
        int n = nums.size();
        int ans = 0, mx = *max_element(nums.begin(), nums.end());
        vector<int> cnt(mx + 2 * k + 1, 0);
        function<void(int)> dfs = [&](int i) { // 第二種：枚舉下一個選哪個
            ans += 1;
            for (int j = i; j < n; j++) { // 枚舉下一個選的是哪個
                int x = nums[j] + k; // 避免索引為負數
                if (cnt[x - k] == 0 && cnt[x + k] == 0) { // 滿足條件
                    cnt[x] += 1;
                    dfs(j + 1);
                    cnt[x] -= 1;
                }
            }
        };
        dfs(0);
        return ans - 1; // 減去空集合
    }
};

解法二：動態規劃(Dynamic Programming) + 乘法原理

首先引入一個前提，兩個整數要相差 $k$ ，則必須滿足 $x \equiv y \pmod{k}$ ，因此可以將陣列中的整數分組，每個組中的整數 $x$ 滿足 $x \equiv i \pmod{k}$ ，其中 $0 \leq i < k$ ，不同組內的選擇不會影響到彼此，因此可以分開計算，再使用乘法原理將結果相乘。

對於每組內的整數，可以使用動態規劃來計算方案數，令 $dp[i]$ 表示前 $i$ 個整數的方案數，注意這裡的 $i$ 是 $1$ -indexed。對於每個整數 $x$ ，都有 $2^{c_x}$ 種選擇，其中 $c_x = cnt[x]$ 表示整數 $x$ 在陣列中出現的次數，在這些選擇中，選 $x$ 的方案數有 $2^{c_x} - 1$ 種，不選 $x$ 的方案數有 $1$ 種。但是若 $x_i$ 和前一個整數 $x_{i-1}$ 相差 $k$ ，則不能同時選擇 $x_i$ 和 $x_{i-1}$ ，因此需要分情況討論：

初始化 $dp[0] = 1$ 和 $dp[1] = 2^{c_{x_1}}$ ，其中 $x_1$ 為第一個整數。
若 $x_i$ $x_{i}$ 和 $x_{i-1}$ $x_{i - 1}$ 相差 $k$ $k$ ，則 $x_i$ $x_{i}$ 的方案數 $dp[i] = dp[i-1] + dp[i-2] \times (2^{c_{x_i}} - 1)$ $d p [i] = d p [i - 1] + d p [i - 2] \times (2^{c_{x_{i}}} - 1)$ 。
- 不選 $x_i$ ，此時可選或可不選 $x_{i-1}$ ，此時方案數為 $dp[i-1]$ 。
- 若選 $x_i$ ，則不能選 $x_{i-1}$ ，而選 $x_i$ 的方案數為 $2^{c_{x_i}} - 1$ ，此時方案數為 $dp[i-2] \times (2^{c_{x_i}} - 1)$ 。
若 $x_i$ 和 $x_{i-1}$ 不相差 $k$ ，則 $x_i$ 的方案數 $dp[i] = dp[i-1] \times 2^{c_{x_i}}$ ，代表可選或可不選 $x_i$ 。

得每組內的方案數為 $dp[m]$ ，其中 $m$ 為組內整數的個數，再將所有組的方案數相乘，減去空集合後即為答案。

複雜度分析

時間複雜度： $\mathcal{O}(n \log n)$ ，看似為 $O(n^2)$ 但其實每個整數只會被分到一組內，因此考慮每組內排序的時間複雜度即可，為 $O(m \log m)$ ，其中 $m$ 為組內整數的最大個數，或可視為 $O(n \log n)$ 。
空間複雜度： $\mathcal{O}(n)$ ，使用雜湊表來記錄每組整數的出現次數，空間複雜度為 $O(n)$ 。

class Solution:
    def beautifulSubsets(self, nums: List[int], k: int) -> int:
        groups = defaultdict(Counter)
        # groups = [Counter() for _ in range(k)]
        for x in nums:
            groups[x % k][x] += 1
        ans = 1
        for _, cnt in groups.items():
        # for cnt in groups:
        #     if not cnt: continue
            m = len(cnt)
            keys = sorted(cnt.keys())
            dp = [0] * (m + 1) # dp[i] 表示前 i 個數字的方案數，1-indexed
            dp[0], dp[1] = 1, 1 << cnt[keys[0]] # 初始化 dp[1] 為 不選和選第一個
            for i in range(1, m):
                # x = keys[i] 的數量有 cnt[x] 個 選 x 的方案有 2^cnt[x] -1 種，不選 x 的方案有 1 種
                x = keys[i]
                if keys[i] - keys[i-1] == k: # 和前一個相差為 k ，則不能同時選 keys[i] 和 keys[i-1]
                    dp[i+1] = dp[i] + dp[i-1] * ((1 << cnt[x]) - 1)
                else:
                    dp[i+1] = dp[i] << cnt[x]
            ans *= dp[m] # 乘法原理
        return ans - 1 # 去除空集合

class Solution {
public:
    int beautifulSubsets(vector<int>& nums, int k) {
        unordered_map<int, map<int, int>> groups;
        for (int x : nums) {
            groups[x % k][x]++;
        }
        int ans = 1;
        for (auto group : groups) { // (key, cnt)
            auto& cnt = group.second;
            int m = cnt.size();
            vector<int> dp(m + 1, 0); // dp[i] 表示前 i 個數字的方案數，1-indexed
            auto it = cnt.begin();
            dp[0] = 1;
            dp[1] = 1 << it++->second; // 初始化 dp[1] 為 不選和選第一個
            for (int i = 1; i < m; i++, it++) {
                int x = it->first;
                if (x - prev(it)->first == k) { // 和前一個相差為 k ，則不能同時選 keys[i] 和 keys[i-1]
                    dp[i + 1] = dp[i] + dp[i - 1] * ((1 << it->second) - 1);
                } else { // 和前一個相差不為 k ，則可以同時選 keys[i] 和 keys[i-1]
                    dp[i + 1] = dp[i] << it->second;
                }
            }
            ans *= dp[m]; // 乘法原理
        }
        return ans - 1; // 去除空集合
    }
};