Find weighted median node in tree

impl Solution {
  pub fn find_median(n: i32, edges: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<i32> {
    let n = n as usize;
    let log = 17usize;

    // Build adjacency list
    let mut adj = vec![vec![]; n];
    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));
    }

    // BFS to set depth, parent, dist from root (0)
    let mut depth = vec![0usize; n];
    let mut parent = vec![vec![0usize; n]; log];
    let mut par_weight = vec![vec![0i64; n]; log];
    let mut dist = vec![0i64; n]; // distance from root
    let mut visited = vec![false; n];
    let mut queue = std::collections::VecDeque::new();
    queue.push_back(0usize);
    visited[0] = true;
    parent[0][0] = 0;
    par_weight[0][0] = 0;

    while let Some(u) = queue.pop_front() {
      for &(v, w) in &adj[u] {
        if !visited[v] {
          visited[v] = true;
          depth[v] = depth[u] + 1;
          parent[0][v] = u;
          par_weight[0][v] = w;
          dist[v] = dist[u] + w;
          queue.push_back(v);
        }
      }
    }

    // Build binary lifting tables
    for k in 1..log {
      for v in 0..n {
        let anc = parent[k - 1][v];
        parent[k][v] = parent[k - 1][anc];
        par_weight[k][v] = par_weight[k - 1][v] + par_weight[k - 1][anc];
      }
    }

    // LCA function
    let lca = |mut u: usize, mut v: usize| -> usize {
      if depth[u] < depth[v] { std::mem::swap(&mut u, &mut v); }
      let diff = depth[u] - depth[v];
      for k in 0..log {
        if (diff >> k) & 1 == 1 {
          u = parent[k][u];
        }
      }
      if u == v { return u; }
      for k in (0..log).rev() {
        if parent[k][u] != parent[k][v] {
          u = parent[k][u];
          v = parent[k][v];
        }
      }
      parent[0][u]
    };

    // Weight from node v up to ancestor anc (exclusive of anc's parent edge)
    // = dist[v] - dist[anc]
    // Walk from u toward v: path is u -> ... -> lca -> ... -> v
    // total_weight = dist[u] + dist[v] - 2*dist[lca]
    // We need to find first node x on path u->v where sum_from_u >= total/2

    let ans: Vec<i32> = queries.iter().map(|q| {
      let (u, v) = (q[0] as usize, q[1] as usize);
      if u == v { return u as i32; }

      let l = lca(u, v);
      let total_w = dist[u] + dist[v] - 2 * dist[l];
      // half threshold: we need cumulative_weight >= ceil(total_w / 2) in integer sense
      // condition is: 2 * cumulative >= total_w
      let half2 = total_w; // we compare 2*cum >= total_w

      // Walk from u toward lca, then from lca toward v
      // Phase 1: u to lca
      // The path from u to lca: go up step by step
      // We need to find first node on path from u upward to lca where dist[u] - dist[node] >= total_w/2
      // i.e., dist[u] - dist[node] >= total_w / 2.0  => 2*(dist[u] - dist[node]) >= total_w

      // Check if median is in u->lca segment:
      // At lca: cum = dist[u] - dist[l]
      // if 2*(dist[u] - dist[l]) >= total_w => median is somewhere in u->lca
      let cu_to_lca = dist[u] - dist[l];
      if 2 * cu_to_lca >= half2 {
        // Binary search: find highest ancestor of u (still below lca) where 2*(dist[u]-dist[anc]) < half2
        // We want: find node x on u->lca path such that 2*(dist[u]-dist[x]) >= half2
        // and 2*(dist[u]-dist[parent(x)]) < half2
        let mut cur = u;
        for k in (0..log).rev() {
          let anc = parent[k][cur];
          // Check if going up 2^k steps still keeps cum < half2
          if depth[anc] >= depth[l] {
            let cum_after = dist[u] - dist[anc];
            if 2 * cum_after < half2 {
              cur = anc;
            }
          }
        }
        // After the loop, parent[0][cur] should be the answer (or cur itself if it was already enough)
        // Actually: cur is the last node where 2*(dist[u]-dist[cur]) < half2
        // So parent[0][cur] is the first node >= half2... but parent[0][cur] might overshoot
        // Let's check: if cur == u and already 2*(dist[u]-dist[u]) = 0 < half2... then move to parent
        let next = parent[0][cur];
        if next == cur {
          // cur is root and the condition is never met on the u-lca segment?
          // Shouldn't happen since we already checked 2*cu_to_lca >= half2
          return l as i32;
        }
        // The answer is `next` (first node going toward lca where cum >= half2)
        // But wait: we need to verify edge from cur to next
        // cum at next = dist[u] - dist[next]
        return next as i32;
      }

      // Median is in lca->v segment (excluding lca itself)
      // cum at lca = dist[u] - dist[l] = cu_to_lca already < half2/2
      // Now walk from lca toward v
      // At node x (descendant of lca on path to v): cum = cu_to_lca + dist[x] - dist[l]
      // We want: 2*(cu_to_lca + dist[x] - dist[l]) >= half2
      // => 2*(dist[x] - dist[l]) >= half2 - 2*cu_to_lca = need
      let need = half2 - 2 * cu_to_lca;
      // Binary lift from v toward lca, jumping UP while destination's cum >= need.
      // This finds the shallowest (closest to lca) node on path lca->v with cum >= need,
      // which is exactly the first node crossing the threshold from lca toward v.
      let mut cur = v;
      for k in (0..log).rev() {
        let anc = parent[k][cur];
        if depth[anc] > depth[l] {
          let cum_at_anc_from_lca = dist[anc] - dist[l];
          if 2 * cum_at_anc_from_lca >= need {
            cur = anc;
          }
        }
      }
      cur as i32
    }).collect();

    ans
  }
}