以下是 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)

 

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐