千问 LeetCode 3915. 距离至少为 K 的交替子序列的最大和 Rust实现
这道题是 LeetCode 3915「距离至少为 K 的交替子序列的最大和」，核心思路是 DP + 值域离散化 + 树状数组（Fenwick Tree）优化，将复杂度从 O(n²) 降到 O(n log n)。
题目理解
给定数组 nums 和整数 k，选一个下标严格递增的子序列，满足：
1. 相邻选中下标之差 ≥ k
2. 选出的值严格交替（小大 或 大小 波动，不能相等）
3. 求最大和
核心思路
1. DP 状态：up[i] 表示以 nums[i] 结尾、最后一步是"递增"（前一个值 < 当前值）的最大和；down[i] 表示以 nums[i] 结尾、最后一步是"递减"的最大和
2. 转移逻辑：
- up[i] = nums[i] + max{down[j]}，其中 j ≤ i-k 且 nums[j] < nums[i]
- down[i] = nums[i] + max{up[j]}，其中 j ≤ i-k 且 nums[j] > nums[i]
3. 延迟激活：只有当 i ≥ k 时，才把 i-k 位置的状态加入树状数组，保证下标距离 ≥ k
4. 树状数组优化：用两棵树状数组分别维护"值小于当前值"和"值大于当前值"的最大 DP 值，查询/更新均为 O(log n)
Rust 实现
use std::cmp::max;
use std::collections::BTreeSet;
struct FenwickTree {
n: usize,
tree: Vec<i64>,
}
impl FenwickTree {
fn new(n: usize) -> Self {
FenwickTree {
n,
tree: vec![i64::MIN / 2; n + 2], // 初始化为极小值
}
}
// 单点取 max 更新
fn update(&mut self, mut idx: usize, val: i64) {
while idx <= self.n {
self.tree[idx] = max(self.tree[idx], val);
idx += idx & idx.wrapping_neg(); // idx += idx & (-idx)
}
}
// 前缀最大值查询 [1, idx]
fn query(&self, mut idx: usize) -> i64 {
let mut res = i64::MIN / 2;
while idx > 0 {
res = max(res, self.tree[idx]);
idx -= idx & idx.wrapping_neg();
}
res
}
}
impl Solution {
pub fn max_alternating_sum(nums: Vec<i32>, k: i32) -> i64 {
let n = nums.len();
let k = k as usize;
// 1. 值域离散化
let mut sorted: Vec<i32> = nums.clone();
sorted.sort();
sorted.dedup();
let m = sorted.len();
// 2. 两棵树状数组
// bit_down：维护 down 值，用于查询"值小于当前值"的最大 down
// bit_up_rev：维护 up 值（倒序坐标），用于查询"值大于当前值"的最大 up
let mut bit_down = FenwickTree::new(m);
let mut bit_up_rev = FenwickTree::new(m);
let mut up = vec![0i64; n];
let mut down = vec![0i64; n];
let mut ans = 0i64;
for i in 0..n {
// 3. 延迟激活：把 i-k 位置的状态加入树状数组
if i >= k {
let prev = i - k;
let prev_rank = sorted.binary_search(&nums[prev]).unwrap() + 1; // 1-based
bit_down.update(prev_rank, down[prev]);
bit_up_rev.update(m - prev_rank + 1, up[prev]); // 倒序映射，后缀变前缀
}
let cur_rank = sorted.binary_search(&nums[i]).unwrap() + 1; // 1-based
// 4. 状态转移
// up[i]：前一个值 < nums[i]，从 bit_down 查询值域 [1, cur_rank-1] 的最大 down
let best_down = bit_down.query(cur_rank - 1);
up[i] = nums[i] as i64 + if best_down <= i64::MIN / 2 { 0 } else { best_down };
// down[i]：前一个值 > nums[i]，从 bit_up_rev 查询值域 [cur_rank+1, m] 的最大 up
let best_up = bit_up_rev.query(m - cur_rank);
down[i] = nums[i] as i64 + if best_up <= i64::MIN / 2 { 0 } else { best_up };
ans = max(ans, max(up[i], down[i]));
}
ans
}
}
关键点解析
- 值域离散化：nums[i] 最大 10⁵，但实际不同值最多 n 个，离散化后压缩到 [1, m]，树状数组大小可控
- 延迟激活：这是处理"下标距离 ≥ k"的关键技巧——遍历时不立即把当前状态加入树状数组，而是等 k 步后再加入，这样查询时自然只看到距离 ≥ k 的前驱状态
- 后缀查询技巧：树状数组天然支持前缀查询，要查"值大于当前值"的最大值，把排名 r 反转为 m - r + 1，就把后缀查询变成了前缀查询
- Rust 特有注意点：idx & (-idx) 在 Rust 中需要用 idx & idx.wrapping_neg() 来避免无符号整数的取负溢出问题；树状数组初始值设为 i64::MIN / 2 防止加法溢出
- 时间复杂度：O(n log n)，空间 O(n)
示例验证
- nums = [5,4,2], k = 2：选下标 [0,2]，值 [5,2]，距离 2-0=2≥k，5>2 严格交替，得分 7 ✅
- nums = [3,5,4,2,4], k = 1：选下标 [0,1,3,4]，值 [3,5,2,4]，3<5>2<4 严格交替，得分 14 ✅
- nums = [5], k = 1：长度为 1 始终有效，得分 5 ✅
需要我把树状数组优化 DP 的通用模板整理出来吗？