Shortest path in a weighted tree

Array Tree Depth-First Search Binary Indexed Tree Segment Tree

42.1% acceptance

Feb 25, 2026

You are given an integer n and an undirected, weighted tree rooted at node 1 with n nodes. edges[i] = [ui, vi, wi] indicates edge from ui to vi with weight wi. queries[i] is either [1, u, v, w'] to update edge weight, or [2, x] to query distance from root. Return answer array for all type-2 queries.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn tree_queries(n: i32, edges: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<i32> {
    use std::collections::HashMap;
    let n = n as usize;
    let mut adj: Vec<Vec<(usize, i64)>> = vec![vec![]; n + 1];
    for e in &edges {
      let (u, v, w) = (e[0] as usize, e[1] as usize, e[2] as i64);
      adj[u].push((v, w));
      adj[v].push((u, w));
    }

    let mut par = vec![0usize; n + 1];
    let mut par_w = vec![0i64; n + 1];
    let mut dist = vec![0i64; n + 1];
    let mut tin = vec![0usize; n + 1];
    let mut tout = vec![0usize; n + 1];
    let mut timer = 0usize;

    let mut stack: Vec<(usize, usize, i64, bool)> = vec![(1, 0, 0, false)];
    while let Some((u, p, d, done)) = stack.pop() {
      if done {
        tout[u] = timer - 1;
      } else {
        par[u] = p;
        dist[u] = d;
        tin[u] = timer;
        timer += 1;
        stack.push((u, p, d, true));
        for &(v, w) in adj[u].iter().rev() {
          if v != p {
            par_w[v] = w;
            stack.push((v, u, d + w, false));
          }
        }
      }
    }

    let mut edge_child: HashMap<(usize, usize), usize> = HashMap::new();
    for e in &edges {
      let u = e[0] as usize;
      let v = e[1] as usize;
      if par[v] == u {
        edge_child.insert((u, v), v);
        edge_child.insert((v, u), v);
      } else {
        edge_child.insert((u, v), u);
        edge_child.insert((v, u), u);
      }
    }

    // BIT: range-add, point-query
    let mut bit = vec![0i64; n + 2];

    let mut result = Vec::new();
    for q in &queries {
      if q[0] == 1 {
        let u = q[1] as usize;
        let v = q[2] as usize;
        let w_new = q[3] as i64;
        let child = *edge_child.get(&(u, v)).unwrap();
        let delta = w_new - par_w[child];
        par_w[child] = w_new;
        let (l, r) = (tin[child], tout[child]);
        Self::bit_add(&mut bit, l, delta);
        Self::bit_add(&mut bit, r + 1, -delta);
      } else {
        let x = q[1] as usize;
        let d = dist[x] + Self::bit_prefix(&bit, tin[x]);
        result.push(d as i32);
      }
    }
    result
  }

  fn bit_add(bit: &mut Vec<i64>, mut i: usize, val: i64) {
    i += 1;
    while i < bit.len() {
      bit[i] += val;
      i += i & i.wrapping_neg();
    }
  }

  fn bit_prefix(bit: &Vec<i64>, mut i: usize) -> i64 {
    i += 1;
    let mut s = 0i64;
    while i > 0 {
      s += bit[i];
      i -= i & i.wrapping_neg();
    }
    s
  }
}