以下是 LeetCode 3721. 最长平衡子数组 II 的 Rust 实现。
---
核心思路
问题转化:子数组中"不同偶数个数 = 不同奇数个数" ⇔ 把每个不同的奇数记为 `+1`,每个不同的偶数记为 `-1`,则子数组平衡 ⇔ 前缀和之差为 0。
关键难点:同一个数字重复出现时,它只贡献一次(去重),但位置变了。需要用线段树维护动态变化的前缀和数组,支持:
1. 区间加(撤销旧贡献 + 添加新贡献)
2. 线段树上二分找最左等于目标值的位置
算法流程:
- 枚举右端点 `i`(1-indexed)
- 若 `nums[i]` 之前出现过,撤销其在旧位置的贡献
- 在当前位置 `i` 添加贡献
- 用线段树查询最早出现相同前缀和的位置 `pos`
- 更新答案 `ans = max(ans, i - pos)`
---
Rust 代码
```rust
use std::collections::HashMap;
/**
* LeetCode 3721. 最长平衡子数组 II
*
* 核心思路:线段树 + 前缀和 + 哈希表
*
* 关键转化:
* - 每个不同的奇数贡献 +1,每个不同的偶数贡献 -1
* - 维护前缀和 now = 不同奇数个数 - 不同偶数个数
* - 子数组 [l, r] 平衡 等价于 prefix[r] - prefix[l-1] = 0
* - 即 prefix[r] == prefix[l-1]
*
* 难点处理:数字重复出现时,需要撤销之前位置的贡献
* - 用线段树维护前缀和数组,支持区间加
* - 用线段树上二分找最左等于目标值的位置
*/
/// 线段树节点
/// 维护区间 [l, r] 的最小值 mn、最大值 mx 和懒标记 lazy
#[derive(Clone, Copy)]
struct Node {
l: usize, // 区间左端点
r: usize, // 区间右端点
mn: i32, // 区间最小值(前缀和)
mx: i32, // 区间最大值(前缀和)
lazy: i32, // 懒标记(区间加)
}
impl Node {
fn new() -> Self {
Node {
l: 0,
r: 0,
mn: 0,
mx: 0,
lazy: 0,
}
}
}
/// 线段树
/// 支持:
/// 1. 区间加
/// 2. 线段树上二分:找最小索引使得前缀和等于 target
struct SegmentTree {
tr: Vec<Node>, // 线段树数组(4倍空间)
}
impl SegmentTree {
/// 创建线段树,区间为 [0, n]
fn new(n: usize) -> Self {
let tr = vec![Node::new(); (n + 1) << 2];
let mut st = SegmentTree { tr };
st.build(1, 0, n);
st
}
/// 建树,初始所有前缀和为 0
fn build(&mut self, u: usize, l: usize, r: usize) {
self.tr[u].l = l;
self.tr[u].r = r;
self.tr[u].mn = 0;
self.tr[u].mx = 0;
self.tr[u].lazy = 0;
if l == r {
return;
}
let mid = (l + r) >> 1;
self.build(u << 1, l, mid);
self.build(u << 1 | 1, mid + 1, r);
}
/// 区间 [l, r] 全部加 v
fn modify(&mut self, u: usize, l: usize, r: usize, v: i32) {
if self.tr[u].l >= l && self.tr[u].r <= r {
self.apply(u, v);
return;
}
self.pushdown(u);
let mid = (self.tr[u].l + self.tr[u].r) >> 1;
if l <= mid {
self.modify(u << 1, l, r, v);
}
if r > mid {
self.modify(u << 1 | 1, l, r, v);
}
self.pushup(u);
}
/// 线段树上二分
/// 找最小索引 pos 使得前缀和 == target
/// 关键观察:如果 target 在 [mn, mx] 范围内,则该区间内一定存在这样的位置
fn query(&mut self, u: usize, target: i32) -> usize {
if self.tr[u].l == self.tr[u].r {
return self.tr[u].l;
}
self.pushdown(u);
let left = u << 1;
let right = u << 1 | 1;
if self.tr[left].mn <= target && target <= self.tr[left].mx {
self.query(left, target)
} else {
self.query(right, target)
}
}
/// 对节点 u 应用区间加 v
fn apply(&mut self, u: usize, v: i32) {
self.tr[u].mn += v;
self.tr[u].mx += v;
self.tr[u].lazy += v;
}
/// 从子节点更新父节点
fn pushup(&mut self, u: usize) {
self.tr[u].mn = self.tr[u << 1].mn.min(self.tr[u << 1 | 1].mn);
self.tr[u].mx = self.tr[u << 1].mx.max(self.tr[u << 1 | 1].mx);
}
/// 下传懒标记
fn pushdown(&mut self, u: usize) {
if self.tr[u].lazy != 0 {
let lazy = self.tr[u].lazy;
self.apply(u << 1, lazy);
self.apply(u << 1 | 1, lazy);
self.tr[u].lazy = 0;
}
}
}
struct Solution;
impl Solution {
pub fn longest_balanced(nums: Vec<i32>) -> i32 {
let n = nums.len();
let mut st = SegmentTree::new(n);
// last[x] = 数值 x 上次出现的位置
let mut last: HashMap<i32, usize> = HashMap::new();
let mut now: i32 = 0; // 当前前缀和
let mut ans: i32 = 0; // 答案
// 枚举子数组右端点(1-indexed)
for i in 1..=n {
let x = nums[i - 1];
// x 的贡献:奇数 +1,偶数 -1
let det = if (x & 1) == 1 { 1 } else { -1 };
// 如果 x 之前出现过,撤销其之前的贡献
if let Some(&pos) = last.get(&x) {
st.modify(1, pos, n, -det);
now -= det;
}
// 添加当前 x 的贡献
last.insert(x, i);
st.modify(1, i, n, det);
now += det;
// 找最早出现相同前缀和的位置
let pos = st.query(1, now);
ans = ans.max((i - pos) as i32);
}
ans
}
}
```
---
复杂度分析
项目 复杂度
每次 modify O(\log n)
每次 query O(\log n)
总时间 O(n \log n)
空间 O(n)(线段树 4n 节点 + 哈希表)
---
示例验证
示例 1:`nums = [2,5,4,3]`
i x det 操作 now pos ans
1 2 -1 modify[1,4]-1 -1 1 0
2 5 +1 modify[2,4]+1 0 0 2
3 4 -1 modify[3,4]-1 -1 1 2
4 3 +1 modify[4,4]+1 0 0 4
最长平衡子数组 `[2,5,4,3]`,长度 4 ✓
示例 2:`nums = [3,2,2,5,4]`
- `i=3` 时 `x=2` 重复出现:撤销 `i=2` 的贡献,在 `i=3` 重新添加
- 最终 `ans = 5`,子数组 `[3,2,2,5,4]` ✓
示例 3:`nums = [1,2,3,2]`
- `i=4` 时 `x=2` 重复出现:撤销 `i=2` 的贡献,在 `i=4` 重新添加
- 最终 `ans = 3`,子数组 `[2,3,2]` ✓
---
下载完整 Rust 文件:[solution_3721.rs](sandbox:///mnt/agents/output/solution_3721.rs)