Kimi LeetCode 3515. 带权树中的最短路径 Rust实现
以下是 LeetCode 3515. 带权树中的最短路径 的完整 Rust 实现。
核心思路
与 Java 版本完全一致:
1. DFS 欧拉序:将每个节点的子树映射到连续区间 `[tin, tout]`
2. 懒标记线段树:区间修改(子树整体加 `delta`)+ 单点查询(节点到根距离)
3. 边权更新:通过 `parent` 数组确定子节点,仅更新子树区间
完整 Rust 代码
```rust
use std::collections::HashMap;
struct Solution;
// ============ Lazy Segment Tree ============
struct LazySegmentTree {
n: usize,
tree: Vec<i64>,
lazy: Vec<i64>,
}
impl LazySegmentTree {
fn new(n: usize) -> Self {
Self {
n,
tree: vec![0; 4 * n],
lazy: vec![0; 4 * n],
}
}
// Range add: add val to every element in [l, r] (inclusive)
fn add_range(&mut self, l: usize, r: usize, val: i64) {
self.add_range_impl(1, 0, self.n - 1, l, r, val);
}
// Point query: get value at index i
fn query(&mut self, i: usize) -> i64 {
self.query_impl(1, 0, self.n - 1, i)
}
fn push(&mut self, id: usize, lo: usize, hi: usize) {
if self.lazy[id] == 0 {
return;
}
self.tree[id] += self.lazy[id];
if lo != hi {
self.lazy[id * 2] += self.lazy[id];
self.lazy[id * 2 + 1] += self.lazy[id];
}
self.lazy[id] = 0;
}
fn add_range_impl(&mut self, id: usize, lo: usize, hi: usize, l: usize, r: usize, val: i64) {
self.push(id, lo, hi);
if r < lo || l > hi {
return;
}
if l <= lo && hi <= r {
self.lazy[id] += val;
self.push(id, lo, hi);
return;
}
let mid = lo + (hi - lo) / 2;
self.add_range_impl(id * 2, lo, mid, l, r, val);
self.add_range_impl(id * 2 + 1, mid + 1, hi, l, r, val);
}
fn query_impl(&mut self, id: usize, lo: usize, hi: usize, i: usize) -> i64 {
self.push(id, lo, hi);
if lo == hi {
return self.tree[id];
}
let mid = lo + (hi - lo) / 2;
if i <= mid {
self.query_impl(id * 2, lo, mid, i)
} else {
self.query_impl(id * 2 + 1, mid + 1, hi, i)
}
}
}
// ============ Graph Edge ============
#[derive(Clone)]
struct Edge {
to: usize,
weight: i32,
}
impl Solution {
pub fn tree_queries(n: i32, edges: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<i32> {
let n = n as usize;
// Build adjacency list
let mut graph: Vec<Vec<Edge>> = vec![Vec::new(); n + 1];
let mut edge_weight: HashMap<(usize, usize), i32> = HashMap::new();
for e in &edges {
let u = e[0] as usize;
let v = e[1] as usize;
let w = e[2];
graph[u].push(Edge { to: v, weight: w });
graph[v].push(Edge { to: u, weight: w });
edge_weight.insert((u.min(v), u.max(v)), w);
}
// Euler tour arrays
let mut tin = vec![0; n + 1];
let mut tout = vec![0; n + 1];
let mut parent = vec![0; n + 1];
let mut dist = vec![0i64; n + 1];
// DFS to compute Euler tour, parent, and initial distances
let mut time: usize = 0;
Self::dfs(&graph, 1, 0, &mut time, &mut tin, &mut tout, &mut dist, &mut parent);
// Build segment tree with initial distances
let mut seg = LazySegmentTree::new(n);
for i in 1..=n {
seg.add_range(tin[i], tin[i], dist[i]);
}
// Process queries
let mut ans: Vec<i32> = Vec::new();
for q in &queries {
if q[0] == 2 {
// Query: shortest path from root to node x
let x = q[1] as usize;
ans.push(seg.query(tin[x]) as i32);
} else {
// Update: change edge (u, v) weight to new_w
let u = q[1] as usize;
let v = q[2] as usize;
let new_w = q[3];
let key = (u.min(v), u.max(v));
let old_w = *edge_weight.get(&key).unwrap();
let delta = (new_w - old_w) as i64;
edge_weight.insert(key, new_w);
// Determine which node is the child (deeper in the tree)
let child = if parent[v] == u { v } else { u };
// Update all nodes in child's subtree
seg.add_range(tin[child], tout[child], delta);
}
}
ans
}
fn dfs(
graph: &Vec<Vec<Edge>>,
u: usize,
prev: usize,
time: &mut usize,
tin: &mut Vec<usize>,
tout: &mut Vec<usize>,
dist: &mut Vec<i64>,
parent: &mut Vec<usize>,
) {
tin[u] = *time;
*time += 1;
for e in &graph[u] {
let v = e.to;
if v == prev {
continue;
}
dist[v] = dist[u] + e.weight as i64;
parent[v] = u;
Self::dfs(graph, v, u, time, tin, tout, dist, parent);
}
tout[u] = *time - 1;
}
}
```
Rust 特有注意点
要点 说明
所有权与借用 `LazySegmentTree` 的方法需要 `&mut self`,因为会修改内部状态
线段树索引 Rust 数组索引为 `usize`,线段树节点从 `1` 开始(`id * 2` 和 `id * 2 + 1`)
类型转换 输入 `i32` 转为 `usize` 用于数组索引,计算时用 `i64` 避免溢出
边权存储 `HashMap<(usize, usize), i32>` 存储无向边,键为有序对 `(min, max)`
递归 DFS 通过 `&mut` 引用传递可变状态(`time`, `tin`, `tout` 等)
复杂度
- 时间复杂度:`O((n + q) log n)`
- 空间复杂度:`O(n)`
下载文件:[LeetCode3515_ShortestPathInWeightedTree.rs](sandbox:///mnt/agents/output/LeetCode3515_ShortestPathInWeightedTree.rs)
更多推荐



所有评论(0)