这是 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 ✓

 

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐