Kimi K2.6思考 LeetCode 3250. 单调数组对的数目 I Rust实现
以下是 LeetCode 3250. 单调数组对的数目 I 的 Rust 实现,采用 动态规划 + 前缀和优化。
核心思路
设 `dp[i][j]` 表示前 `i` 个元素,且 `arr1[i] = j` 时的方案数。约束条件:
1. `arr1` 非递减:`arr1[i-1] <= arr1[i]` → `prev_j <= j`
2. `arr2` 非递增:`arr2[i-1] >= arr2[i]` → `nums[i-1] - prev_j >= nums[i] - j` → `prev_j <= j + nums[i-1] - nums[i]`
合并:`prev_j <= min(j, j + nums[i-1] - nums[i])`
利用前缀和将转移优化到 O(1),整体复杂度 O(n × m),其中 `m = max(nums) <= 50`。
Rust 代码
```rust
impl Solution {
pub fn count_of_pairs(nums: Vec<i32>) -> i32 {
const MOD: i64 = 1_000_000_007;
let n = nums.len();
if n == 0 {
return 0;
}
let max_num = *nums.iter().max().unwrap() as usize;
// dp[j] 表示当前位置,arr1[i] = j 的方案数
let mut dp = vec![0i64; max_num + 1];
// 初始化:i = 0 时,arr1[0] 可以是 0..=nums[0]
for j in 0..=nums[0] as usize {
dp[j] = 1;
}
for i in 1..n {
let mut new_dp = vec![0i64; max_num + 1];
// 前缀和数组:prefix[j] = sum(dp[0..=j])
let mut prefix = vec![0i64; max_num + 1];
prefix[0] = dp[0];
for j in 1..=max_num {
prefix[j] = (prefix[j - 1] + dp[j]) % MOD;
}
let curr_num = nums[i] as i32;
let prev_num = nums[i - 1] as i32;
for j in 0..=curr_num as usize {
// prev_j 的上界:min(j, j + nums[i-1] - nums[i])
let limit = (j as i32).min(j as i32 + prev_num - curr_num) as usize;
if limit as i32 > max_num as i32 || limit as i32 > j as i32 + prev_num - curr_num {
// 处理负数情况
new_dp[j] = 0;
} else if limit > max_num {
new_dp[j] = 0;
} else {
new_dp[j] = prefix[limit];
}
}
dp = new_dp;
}
let ans: i64 = dp[..=nums[n - 1] as usize].iter().sum::<i64>() % MOD;
ans as i32
}
}
```
更简洁的版本(处理负数边界)
```rust
impl Solution {
pub fn count_of_pairs(nums: Vec<i32>) -> i32 {
const MOD: i64 = 1_000_000_007;
let n = nums.len();
let max_num = *nums.iter().max().unwrap() as usize;
let mut dp = vec![0i64; max_num + 1];
// 初始化
for j in 0..=nums[0] as usize {
dp[j] = 1;
}
for i in 1..n {
let mut new_dp = vec![0i64; max_num + 1];
let mut prefix = vec![0i64; max_num + 1];
prefix[0] = dp[0];
for j in 1..=max_num {
prefix[j] = (prefix[j - 1] + dp[j]) % MOD;
}
let diff = nums[i - 1] - nums[i]; // nums[i-1] - nums[i]
for j in 0..=nums[i] as usize {
// prev_j <= j 且 prev_j <= j + diff
// 即 prev_j <= min(j, j + diff) = j + min(0, diff)
let limit = if diff < 0 {
// diff < 0 时,j + diff < j,所以 limit = j + diff
if j as i32 + diff < 0 {
// 没有合法的 prev_j
continue;
}
(j as i32 + diff) as usize
} else {
// diff >= 0 时,j <= j + diff,所以 limit = j
j
};
if limit <= max_num {
new_dp[j] = prefix[limit];
}
}
dp = new_dp;
}
let ans = dp[..=nums[n - 1] as usize].iter().fold(0i64, |acc, &x| (acc + x) % MOD);
ans as i32
}
}
```
最简洁的版本(推荐)
```rust
impl Solution {
pub fn count_of_pairs(nums: Vec<i32>) -> i32 {
const MOD: i64 = 1_000_000_007;
let max_num = *nums.iter().max().unwrap() as usize;
let mut dp = vec![0i64; max_num + 1];
for j in 0..=nums[0] as usize {
dp[j] = 1;
}
for i in 1..nums.len() {
let mut prefix = vec![0i64; max_num + 1];
prefix[0] = dp[0];
for j in 1..=max_num {
prefix[j] = (prefix[j - 1] + dp[j]) % MOD;
}
let mut new_dp = vec![0i64; max_num + 1];
let diff = nums[i - 1] - nums[i];
for j in 0..=nums[i] as usize {
let limit = (j as i32).min(j as i32 + diff);
if limit >= 0 {
new_dp[j] = prefix[limit as usize];
}
}
dp = new_dp;
}
dp[..=nums[nums.len() - 1] as usize]
.iter()
.fold(0i64, |a, b| (a + b) % MOD) as i32
}
}
```
复杂度分析
维度 复杂度
时间 O(n × m),`n <= 2000`,`m = max(nums) <= 50`
空间 O(m),滚动数组优化
关键点
- 前缀和优化:`prefix[j] = sum(dp[0..=j])`,将区间求和降到 O(1)
- 负数处理:`limit = min(j, j + diff)` 可能为负,此时无合法前驱状态,直接跳过
- 类型转换:Rust 中 `usize` 不能为负,需用 `i32` 计算 `limit` 后再判断是否为非负数
- 取模:全程使用 `i64` 防止中间溢出,最后转回 `i32`
更多推荐

所有评论(0)