这是 LeetCode 3883 的 Rust 实现,基于动态规划 + 双指针优化的思路。
解题思路
1. 预处理:枚举 `0..=5000` 所有数字,按数位和分组。最大数位和为 `4+9+9+9=31`(数字4999),所以最多32个组。
2. 动态规划:`f[x]` 表示以数字 `x` 结尾的有效非递减数组数量。设虚拟起点 `f[0] = 1`(前一个数位和为0,数字为0)。
3. 双指针转移:对于当前数位和 `cur`,遍历该组所有数字 `x`。用双指针在上一组(数位和 `pre`)中累加所有 `<= x` 的 `f` 值,即为以 `x` 结尾的方案数。
4. 取模:结果对 `10^9 + 7` 取模。
```rust
use std::sync::OnceLock;
const MOD: i64 = 1_000_000_007;
const MX: usize = 5001;
const MAX_DIGIT_SUM: usize = 31; // 4999 的数位和最大为 31
/// 预计算:按数位和分组数字 0..5000
fn get_sum_to_nums() -> &'static Vec<Vec<usize>> {
static SUM_TO_NUMS: OnceLock<Vec<Vec<usize>>> = OnceLock::new();
SUM_TO_NUMS.get_or_init(|| {
let mut sum_to_nums: Vec<Vec<usize>> = vec![Vec::new(); MAX_DIGIT_SUM + 1];
let mut dig_sum = vec![0usize; MX];
for x in 0..MX {
if x > 0 {
dig_sum[x] = dig_sum[x / 10] + (x % 10);
}
sum_to_nums[dig_sum[x]].push(x);
}
sum_to_nums
})
}
struct Solution;
impl Solution {
pub fn count_arrays(digit_sum: Vec<i32>) -> i32 {
let sum_to_nums = get_sum_to_nums();
let mut f: Vec<i64> = vec![0; MX];
f[0] = 1; // 以 0 结尾的方案数为 1(虚起点)
let mut pre: usize = 0; // 上一个数位和,初始为 0
for &cur in &digit_sum {
let cur = cur as usize;
if cur > MAX_DIGIT_SUM {
return 0;
}
let a = &sum_to_nums[pre]; // 前一个数位和对应的数字列表
let b = &sum_to_nums[cur]; // 当前数位和对应的数字列表
let mut j: usize = 0;
let m = a.len();
let mut sum: i64 = 0;
for &x in b {
// 累加所有满足 a[j] <= x 的 f[a[j]]
while j < m && a[j] <= x {
sum += f[a[j]];
j += 1;
}
f[x] = sum % MOD;
}
pre = cur; // 更新前一个数位和
}
let mut ans: i64 = 0;
for &x in &sum_to_nums[pre] {
ans += f[x];
}
(ans % MOD) as i32
}
}
fn main() {
// 示例测试
let tests = vec![
(vec![25, 1], 6),
(vec![1], 4),
(vec![2, 49, 23], 0),
];
for (digit_sum, expected) in tests {
let result = Solution::count_arrays(digit_sum.clone());
println!("digitSum = {:?}, result = {}, expected = {}", digit_sum, result, expected);
assert_eq!(result, expected);
}
}
```
复杂度分析
- 时间复杂度:预处理 `O(5000)`;DP 阶段每步双指针总扫描量不超过两组长度之和,总体 `O(n × L_max)`,其中 `L_max` 是单组最大长度(实际很小)。
- 空间复杂度:`O(5000)`,用于 `f` 数组和预计算的分组。