Longest special path ii

impl Solution {
  pub fn longest_special_path(edges: Vec<Vec<i32>>, nums: Vec<i32>) -> Vec<i32> {
    let n = nums.len();
    let mut adj: Vec<Vec<(usize, i64)>> = 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));
    }
    let mut best_len = 0i64;
    let mut best_nodes = 1i32;
    // prefix[d] = total edge weight from root (depth 0) down to the node at depth d.
    // Maintained as a stack: pushed when descending, popped when backtracking.
    let mut prefix = vec![0i64];
    // occ[val] = sorted list of depths at which `val` appears on the current root→node path.
    let mut occ: std::collections::HashMap<i32, Vec<usize>> = std::collections::HashMap::new();
    // Invariant: before entering each node, the window [min_start, d-1] has at most 1 duplicate value.
    // dup_val = -1 means no dup; otherwise it holds the one duplicated value.
    let mut dup_val: i32 = -1;
    let mut min_start: usize = 0;

    fn dfs(
      node: usize, parent: usize, depth: usize,
      adj: &Vec<Vec<(usize, i64)>>, nums: &[i32],
      prefix: &mut Vec<i64>,
      occ: &mut std::collections::HashMap<i32, Vec<usize>>,
      dup_val: &mut i32,
      min_start: &mut usize,
      best_len: &mut i64,
      best_nodes: &mut i32,
    ) {
      let val = nums[node];
      let d = depth;
      let saved_min = *min_start;
      let saved_dup = *dup_val;

      occ.entry(val).or_default().push(d);

      // Compute how many times `val` appears in the current window [min_start, d].
      // Also capture the "first in window" index and second-to-last occurrence for reuse.
      let (win_occ, first_win_idx) = {
        let o = occ.get(&val).unwrap();
        let fwi = o.partition_point(|&x| x < *min_start);
        (o.len() - fwi, fwi)
      };

      if win_occ >= 3 {
        // `val` was already the single dup (invariant); it now appears 3 times.
        // Advance min_start past its first in-window occurrence so it stays at exactly 2.
        let first_in_win = occ.get(&val).unwrap()[first_win_idx];
        *min_start = first_in_win + 1;
        // dup_val stays = val (still 2 occurrences in new window)
      } else if win_occ == 2 {
        // `val` just became a duplicate.
        // fv = the occurrence that was already in the window before this push = o[-2]
        let fv = {
          let o = occ.get(&val).unwrap();
          o[o.len() - 2]
        };

        if *dup_val == -1 {
          // No prior dup; val becomes the single dup.
          *dup_val = val;
        } else {
          // There is already another dup `u`. Advance min_start to leave at most 1 dup.
          // fu = u's first occurrence in the current window.
          let u = *dup_val;
          let fu = {
            let ou = occ.get(&u).unwrap();
            let fu_idx = ou.partition_point(|&x| x < *min_start);
            ou[fu_idx]
          };

          if fu < fv {
            // Advance past u's first; u loses its duplicate status, val keeps it.
            *min_start = fu + 1;
            *dup_val = val;
          } else if fv < fu {
            // Advance past val's first; val loses its duplicate status, u keeps it.
            *min_start = fv + 1;
            // dup_val stays = u
          } else {
            // fu == fv: both lose their first occurrence; neither remains a dup.
            *min_start = fu + 1;
            *dup_val = -1;
          }
        }
      }

      // Window is now valid. Compute length and node count.
      let len = prefix[d] - prefix[*min_start];
      let nodes = (d + 1 - *min_start) as i32;
      if len > *best_len || (len == *best_len && nodes < *best_nodes) {
        *best_len = len;
        *best_nodes = nodes;
      }

      for &(nb, w) in &adj[node] {
        if nb == parent { continue; }
        prefix.push(prefix[d] + w);
        dfs(nb, node, d + 1, adj, nums, prefix, occ, dup_val, min_start, best_len, best_nodes);
        prefix.pop();
      }

      // Backtrack: restore state
      occ.get_mut(&val).unwrap().pop();
      *min_start = saved_min;
      *dup_val = saved_dup;
    }

    dfs(0, usize::MAX, 0, &adj, &nums, &mut prefix, &mut occ, &mut dup_val, &mut min_start, &mut best_len, &mut best_nodes);
    vec![best_len as i32, best_nodes]
  }
}