根据之前搜索到的题目描述和解题思路,我直接为你编写 Go 实现。
LeetCode 3883 - Count Non Decreasing Arrays With Given Digit Sums (Go)
解题思路
1. 预处理:枚举 `[0, 5000]` 中每个数字,按其数位和分组
2. 动态规划 + 前缀和优化:`dp[j]` 表示前一个位置选择第 `j` 个候选数字的方案数,用前缀和加速"非递减"约束的转移
```go
package main
import (
"sort"
)
const MOD = 1_000_000_007
func countArrays(digitSum []int) int {
n := len(digitSum)
// 预处理:按数位和分组,收集 [0, 5000] 中的数字
// digitSum[i] <= 50,所以只需要 0~50 的分组
candidates := make([][]int, 51)
for num := 0; num <= 5000; num++ {
s := digitSumOf(num)
if s <= 50 {
candidates[s] = append(candidates[s], num)
}
}
// 如果某个位置没有候选数字,直接返回 0
for _, s := range digitSum {
if len(candidates[s]) == 0 {
return 0
}
}
// 确保每个分组都是有序的(由于按 num 顺序添加,已经有序)
for i := range candidates {
sort.Ints(candidates[i])
}
// 初始化第一个位置:每个候选数字都是一种方案
first := candidates[digitSum[0]]
dp := make([]int, len(first))
for i := range dp {
dp[i] = 1
}
// 处理后续每个位置
for i := 1; i < n; i++ {
prev := candidates[digitSum[i-1]]
curr := candidates[digitSum[i]]
// 计算 dp 的前缀和
prefix := make([]int, len(prev))
prefix[0] = dp[0]
for j := 1; j < len(prev); j++ {
prefix[j] = (prefix[j-1] + dp[j]) % MOD
}
newDp := make([]int, len(curr))
// 双指针:对于当前每个候选,找到所有 <= 它的 prev 候选
p := 0
for j := 0; j < len(curr); j++ {
currVal := curr[j]
for p < len(prev) && prev[p] <= currVal {
p++
}
// prev[0..p-1] 都 <= currVal
if p > 0 {
newDp[j] = prefix[p-1]
}
}
dp = newDp
}
// 统计所有方案
ans := 0
for _, v := range dp {
ans = (ans + v) % MOD
}
return ans
}
// 计算数字的数位和
func digitSumOf(num int) int {
sum := 0
for num > 0 {
sum += num % 10
num /= 10
}
return sum
}
```
关键点说明
要点 说明
预处理 枚举 `0~5000`,按数位和 `0~50` 分组,每组最多约 100 个数字
双指针优化 两个候选列表都是有序的,用双指针在 O(\|prev\| + \|curr\|) 内完成转移
前缀和 `prefix[j]` 表示 `dp[0..j]` 的和,快速求"所有合法前驱的方案数之和"
空间优化 只保留一维 DP,空间复杂度 O(m),m 为候选数字数量
复杂度
- 时间:O(n × m),n ≤ 1000,m 为每组候选数数量
- 空间:O(m)
示例验证
- `digitSum = [25, 1]` → 输出 `6`(799/889/898/979/988/997 后面接 1000)
- `digitSum = [1]` → 输出 `4`(1, 10, 100, 1000)
- `digitSum = [2, 49, 23]` → 输出 `0`(49 在 [0,5000] 内无数位和为 49 的数字)