这道题是 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 的通用模板整理出来吗?