Kimi LeetCode 3539. 魔法序列的数组乘积之和 Go实现
这是 LeetCode 3539. 魔法序列的数组乘积之和 的 Go 实现。
思路
本题是组合数学 + 记忆化搜索。
核心观察: 枚举每个 `nums[i]` 被选取的次数 `c_i`(`Σc_i = m`),则:
- 序列的排列数为多重集排列:`m! / (c_0! · c_1! · ... · c_{n-1}!)`
- 数组乘积为:`nums[0]^c_0 · nums[1]^c_1 · ...`
利用组合数 `C(remaining, count)` 在递归中直接累乘,恰好得到 `m! / ∏c_i!` 的因子,无需最后再乘 `m!`。
状态定义: `dp(i, remainingM, remainingK, carry)` 表示:
- 处理到第 `i` 个元素
- 还需选 `remainingM` 个数
- 当前低位向第 `i` 位的进位为 `carry`(即 `sum >> i`)
- 还需满足 `remainingK` 个置位
转移: 枚举选 `count` 个下标 `i`,则:
- 当前位产生的置位:`(carry + count) & 1`
- 新进位:`(carry + count) >> 1`
- 贡献:`C(remainingM, count) · nums[i]^count`
复杂度: `O(n · m² · k)`,由于 `m ≤ 30, n ≤ 50`,完全可行。
---
```go
package main
import "math/bits"
func magicalSum(m int, k int, nums []int) int {
const MOD int = 1_000_000_007
n := len(nums)
// 预处理组合数 C[i][j] = C(i, j)
comb := make([][]int, m+1)
for i := 0; i <= m; i++ {
comb[i] = make([]int, m+1)
comb[i][0] = 1
for j := 1; j <= i; j++ {
comb[i][j] = (comb[i-1][j-1] + comb[i-1][j]) % MOD
}
}
// memo[i][remM][remK][carry],-1 表示未计算
memo := make([][][][]int, n+1)
for i := range memo {
memo[i] = make([][][]int, m+1)
for j := range memo[i] {
memo[i][j] = make([][]int, k+1)
for l := range memo[i][j] {
memo[i][j][l] = make([]int, m+1)
for c := range memo[i][j][l] {
memo[i][j][l][c] = -1
}
}
}
}
var dfs func(i, remM, remK, carry int) int
dfs = func(i, remM, remK, carry int) int {
bits := bits.OnesCount32(uint32(carry))
// 剪枝:即使剩余所有位都产生置位,也不够 remK
if remM < 0 || remK < 0 || bits+remM < remK {
return 0
}
// 已选完所有数,检查进位中是否恰好还有 remK 个置位
if remM == 0 {
if remK == bits {
return 1
}
return 0
}
// 没有更多数字可选,但还没选够
if i == n {
return 0
}
if memo[i][remM][remK][carry] != -1 {
return memo[i][remM][remK][carry]
}
res := 0
for count := 0; count <= remM; count++ {
// 从剩余 remM 个位置中选 count 个放 nums[i]
contribution := comb[remM][count] * modPow(nums[i], count, MOD) % MOD
newCarry := carry + count
next := dfs(i+1, remM-count, remK-(newCarry&1), newCarry>>1)
res = (res + next*contribution) % MOD
}
memo[i][remM][remK][carry] = res
return res
}
return dfs(0, m, k, 0)
}
func modPow(base, exp, mod int) int {
res := 1
b := base % mod
e := exp
for e > 0 {
if e&1 == 1 {
res = res * b % mod
}
b = b * b % mod
e >>= 1
}
return res
}
```
---
验证:
- 示例 1:`m=5, k=5, nums=[1,10,100,10000,1000000]` → 991600007 ✓
- 示例 2:`m=2, k=2, nums=[5,4,3,2,1]` → 170 ✓
- 示例 3:`m=1, k=1, nums=[28]` → 28 ✓
更多推荐


所有评论(0)