Checking existence of edge length limited paths ii

struct DistanceLimitedPathsExist {
  parent: Vec<Vec<i32>>,
  max_weight: Vec<Vec<i32>>,
  depth: Vec<i32>,
  component: Vec<i32>,
}


/** 
 * `&self` means the method takes an immutable reference.
 * If you need a mutable reference, change it to `&mut self` instead.
 */
impl DistanceLimitedPathsExist {

  fn new(n: i32, edge_list: Vec<Vec<i32>>) -> Self {
    let n = n as usize;
    let log_n = 15;
    let mut edges: Vec<(i32, usize, usize)> = edge_list.iter()
      .map(|e| (e[2], e[0] as usize, e[1] as usize))
      .collect();
    edges.sort();

    let mut uf: Vec<usize> = (0..n).collect();
    let mut rank = vec![0usize; n];

    fn find(uf: &mut Vec<usize>, x: usize) -> usize {
      if uf[x] != x { uf[x] = find(uf, uf[x]); }
      uf[x]
    }

    let mut adj = vec![vec![]; n];
    for &(w, u, v) in &edges {
      let ru = find(&mut uf, u);
      let rv = find(&mut uf, v);
      if ru != rv {
        adj[u].push((v, w));
        adj[v].push((u, w));
        if rank[ru] < rank[rv] { uf[ru] = rv; }
        else if rank[ru] > rank[rv] { uf[rv] = ru; }
        else { uf[rv] = ru; rank[ru] += 1; }
      }
    }

    let mut component = vec![-1i32; n];
    let mut parent_table = vec![vec![-1i32; n]; log_n];
    let mut weight_table = vec![vec![0i32; n]; log_n];
    let mut depth = vec![0i32; n];
    let mut visited = vec![false; n];
    let mut comp_id = 0i32;

    for start in 0..n {
      if visited[start] { continue; }
      visited[start] = true;
      component[start] = comp_id;
      let mut queue = std::collections::VecDeque::new();
      queue.push_back(start);
      while let Some(u) = queue.pop_front() {
        for &(v, w) in &adj[u] {
          if !visited[v] {
            visited[v] = true;
            component[v] = comp_id;
            depth[v] = depth[u] + 1;
            parent_table[0][v] = u as i32;
            weight_table[0][v] = w;
            queue.push_back(v);
          }
        }
      }
      comp_id += 1;
    }

    for k in 1..log_n {
      for i in 0..n {
        let p = parent_table[k - 1][i];
        if p != -1 {
          parent_table[k][i] = parent_table[k - 1][p as usize];
          weight_table[k][i] = weight_table[k - 1][i].max(weight_table[k - 1][p as usize]);
        }
      }
    }

    DistanceLimitedPathsExist { parent: parent_table, max_weight: weight_table, depth, component }
  }
  
  fn query(&self, p: i32, q: i32, limit: i32) -> bool {
    let (mut a, mut b) = (p as usize, q as usize);
    if self.component[a] != self.component[b] { return false; }

    let log_n = self.parent.len();
    let mut max_w = 0i32;

    if self.depth[a] < self.depth[b] { std::mem::swap(&mut a, &mut b); }

    let diff = self.depth[a] - self.depth[b];
    for k in 0..log_n {
      if diff & (1 << k) != 0 {
        max_w = max_w.max(self.max_weight[k][a]);
        a = self.parent[k][a] as usize;
      }
    }

    if a == b { return max_w < limit; }

    for k in (0..log_n).rev() {
      if self.parent[k][a] != self.parent[k][b] {
        max_w = max_w.max(self.max_weight[k][a]);
        max_w = max_w.max(self.max_weight[k][b]);
        a = self.parent[k][a] as usize;
        b = self.parent[k][b] as usize;
      }
    }

    max_w = max_w.max(self.max_weight[0][a]);
    max_w = max_w.max(self.max_weight[0][b]);
    max_w < limit
  }
}

/*
 * Your DistanceLimitedPathsExist object will be instantiated and called as such:
 * let obj = DistanceLimitedPathsExist::new(n, edgeList);
 * let ret_1: bool = obj.query(p, q, limit);
 */