feat: add solutions to lc problem: No.3539 (#4783)

Author: yanglbme

solution/3500-3599/3539.Find Sum of Array Product of Magical Sequences/README.md

@@ -94,32 +94,297 @@ tags:
 
 <!-- solution:start -->
 
-### 方法一
+### 方法一：组合数学 + 记忆化搜索
+
+我们设计一个函数 $\text{dfs}(i, j, k, st)$，表示当前处理到数组 $\textit{nums}$ 的第 $i$ 个元素，当前还需要从剩余的 $j$ 个位置中选择数字填入魔法序列，当前还需要满足二进制形式有 $k$ 个置位，当前上一位的进位为 $st$ 的方案数。那么答案为 $\text{dfs}(0, m, k, 0)$。
+
+函数 $\text{dfs}(i, j, k, st)$ 的执行流程如下：
+
+如果 $k < 0$ 或者 $i = n$ 且 $j > 0$，说明当前方案不可行，返回 $0$。
+
+如果 $i = n$，说明已经处理完数组 $\textit{nums}$，我们需要检查当前进位 $st$ 中是否还有置位，如果有则需要减少 $k$。如果此时 $k = 0$，说明当前方案可行，返回 $1$，否则返回 $0$。
+
+否则，我们枚举在位置 $i$ 选择 $t$ 个数字填入魔法序列（$0 \leq t \leq j$），将 $t$ 个数字填入魔法序列的方案数为 $\binom{j}{t}$，数组乘积为 $\textit{nums}[i]^t$，更新进位为 $(t + st) >> 1$，更新需要满足的置位数为 $k - ((t + st) \& 1)$，递归调用 $\text{dfs}(i + 1, j - t, k - ((t + st) \& 1), (t + st) >> 1)$。将所有 $t$ 的方案数累加即为 $\text{dfs}(i, j, k, st)$。
+
+为了高效计算组合数 $\binom{m}{n}$，我们预处理阶乘数组 $f$ 和阶乘的逆元数组 $g$，其中 $f[i] = i! \mod (10^9 + 7)$，$g[i] = (i!)^{-1} \mod (10^9 + 7)$。则 $\binom{m}{n} = f[m] \cdot g[n] \cdot g[m - n] \mod (10^9 + 7)$。
+
+时间复杂度 $O(n \cdot m^3 \cdot k)$，空间复杂度 $O(n \cdot m^2 \cdot k)$，其中 $n$ 是数组 $\textit{nums}$ 的长度，而 $m$ 和 $k$ 分别是题目中的参数。
 
 <!-- tabs:start -->
 
 #### Python3
 
 ```python
-
+mx = 30
+mod = 10**9 + 7
+f = [1] + [0] * mx
+g = [1] + [0] * mx
+
+for i in range(1, mx + 1):
+    f[i] = f[i - 1] * i % mod
+    g[i] = pow(f[i], mod - 2, mod)
+
+
+def comb(m: int, n: int) -> int:
+    return f[m] * g[n] * g[m - n] % mod
+
+
+class Solution:
+    def magicalSum(self, m: int, k: int, nums: List[int]) -> int:
+        @cache
+        def dfs(i: int, j: int, k: int, st: int) -> int:
+            if k < 0 or (i == len(nums) and j > 0):
+                return 0
+            if i == len(nums):
+                while st:
+                    k -= st & 1
+                    st >>= 1
+                return int(k == 0)
+            res = 0
+            for t in range(j + 1):
+                nt = t + st
+                p = pow(nums[i], t, mod)
+                nk = k - (nt & 1)
+                res += comb(j, t) * p * dfs(i + 1, j - t, nk, nt >> 1)
+                res %= mod
+            return res
+
+        ans = dfs(0, m, k, 0)
+        dfs.cache_clear()
+        return ans
 ```
 
 #### Java
 
 ```java
-
+class Solution {
+    static final int N = 31;
+    static final long MOD = 1_000_000_007L;
+    private static final long[] f = new long[N];
+    private static final long[] g = new long[N];
+    private Long[][][][] dp;
+
+    static {
+        f[0] = 1;
+        g[0] = 1;
+        for (int i = 1; i < N; ++i) {
+            f[i] = f[i - 1] * i % MOD;
+            g[i] = qpow(f[i], MOD - 2);
+        }
+    }
+
+    public static long qpow(long a, long k) {
+        long res = 1;
+        while (k != 0) {
+            if ((k & 1) == 1) {
+                res = res * a % MOD;
+            }
+            a = a * a % MOD;
+            k >>= 1;
+        }
+        return res;
+    }
+
+    public static long comb(int m, int n) {
+        return f[m] * g[n] % MOD * g[m - n] % MOD;
+    }
+
+    public int magicalSum(int m, int k, int[] nums) {
+        int n = nums.length;
+        dp = new Long[n + 1][m + 1][k + 1][N];
+        long ans = dfs(0, m, k, 0, nums);
+        return (int) ans;
+    }
+
+    private long dfs(int i, int j, int k, int st, int[] nums) {
+        if (k < 0 || (i == nums.length && j > 0)) {
+            return 0;
+        }
+        if (i == nums.length) {
+            while (st > 0) {
+                k -= (st & 1);
+                st >>= 1;
+            }
+            return k == 0 ? 1 : 0;
+        }
+
+        if (dp[i][j][k][st] != null) {
+            return dp[i][j][k][st];
+        }
+
+        long res = 0;
+        for (int t = 0; t <= j; t++) {
+            int nt = t + st;
+            int nk = k - (nt & 1);
+            long p = qpow(nums[i], t);
+            long tmp = comb(j, t) * p % MOD * dfs(i + 1, j - t, nk, nt >> 1, nums) % MOD;
+            res = (res + tmp) % MOD;
+        }
+
+        return dp[i][j][k][st] = res;
+    }
+}
 ```
 
 #### C++
 
 ```cpp
-
+const int N = 31;
+const long long MOD = 1'000'000'007;
+
+long long f[N], g[N];
+
+long long qpow(long long a, long long k) {
+    long long res = 1;
+    while (k) {
+        if (k & 1) res = res * a % MOD;
+        a = a * a % MOD;
+        k >>= 1;
+    }
+    return res;
+}
+
+int init = []() {
+    f[0] = g[0] = 1;
+    for (int i = 1; i < N; ++i) {
+        f[i] = f[i - 1] * i % MOD;
+        g[i] = qpow(f[i], MOD - 2);
+    }
+    return 0;
+}();
+
+long long comb(int m, int n) {
+    return f[m] * g[n] % MOD * g[m - n] % MOD;
+}
+
+class Solution {
+    vector<vector<vector<vector<long long>>>> dp;
+
+    long long dfs(int i, int j, int k, int st) {
+        if (k < 0 || (i == nums.size() && j > 0)) {
+            return 0;
+        }
+        if (i == nums.size()) {
+            while (st > 0) {
+                k -= (st & 1);
+                st >>= 1;
+            }
+            return k == 0 ? 1 : 0;
+        }
+
+        long long& res = dp[i][j][k][st];
+        if (res != -1) {
+            return res;
+        }
+
+        res = 0;
+        for (int t = 0; t <= j; ++t) {
+            int nt = t + st;
+            int nk = k - (nt & 1);
+            long long p = qpow(nums[i], t);
+            long long tmp = comb(j, t) * p % MOD * dfs(i + 1, j - t, nk, nt >> 1) % MOD;
+            res = (res + tmp) % MOD;
+        }
+        return res;
+    }
+
+public:
+    int magicalSum(int m, int k, vector<int>& nums) {
+        int n = nums.size();
+        this->nums = nums;
+        dp.assign(n + 1, vector<vector<vector<long long>>>(m + 1, vector<vector<long long>>(k + 1, vector<long long>(N, -1))));
+        return dfs(0, m, k, 0);
+    }
+
+private:
+    vector<int> nums;
+};
 ```
 
 #### Go
 
 ```go
-
+const N = 31
+const MOD = 1_000_000_007
+
+var f [N]int64
+var g [N]int64
+
+func init() {
+	f[0], g[0] = 1, 1
+	for i := 1; i < N; i++ {
+		f[i] = f[i-1] * int64(i) % MOD
+		g[i] = qpow(f[i], MOD-2)
+	}
+}
+
+func qpow(a, k int64) int64 {
+	res := int64(1)
+	for k > 0 {
+		if k&1 == 1 {
+			res = res * a % MOD
+		}
+		a = a * a % MOD
+		k >>= 1
+	}
+	return res
+}
+
+func comb(m, n int) int64 {
+	if n < 0 || n > m {
+		return 0
+	}
+	return f[m] * g[n] % MOD * g[m-n] % MOD
+}
+
+func magicalSum(m int, k int, nums []int) int {
+	n := len(nums)
+	dp := make([][][][]int64, n+1)
+	for i := 0; i <= n; i++ {
+		dp[i] = make([][][]int64, m+1)
+		for j := 0; j <= m; j++ {
+			dp[i][j] = make([][]int64, k+1)
+			for l := 0; l <= k; l++ {
+				dp[i][j][l] = make([]int64, N)
+				for s := 0; s < N; s++ {
+					dp[i][j][l][s] = -1
+				}
+			}
+		}
+	}
+
+	var dfs func(i, j, k, st int) int64
+	dfs = func(i, j, k, st int) int64 {
+		if k < 0 || (i == n && j > 0) {
+			return 0
+		}
+		if i == n {
+			for st > 0 {
+				k -= st & 1
+				st >>= 1
+			}
+			if k == 0 {
+				return 1
+			}
+			return 0
+		}
+		if dp[i][j][k][st] != -1 {
+			return dp[i][j][k][st]
+		}
+		res := int64(0)
+		for t := 0; t <= j; t++ {
+			nt := t + st
+			nk := k - (nt & 1)
+			p := qpow(int64(nums[i]), int64(t))
+			tmp := comb(j, t) * p % MOD * dfs(i+1, j-t, nk, nt>>1) % MOD
+			res = (res + tmp) % MOD
+		}
+		dp[i][j][k][st] = res
+		return res
+	}
+
+	return int(dfs(0, m, k, 0))
+}
 ```
 
 <!-- tabs:end -->