Subtree inversion sum

impl Solution {
  pub fn subtree_inversion_sum(edges: Vec<Vec<i32>>, nums: Vec<i32>, k: i32) -> i64 {
    let n = nums.len();
    let k = k as usize;
    
    // Build adjacency list
    let mut adj = vec![vec![]; n];
    for e in &edges {
      let (u, v) = (e[0] as usize, e[1] as usize);
      adj[u].push(v);
      adj[v].push(u);
    }
    
    // DFS with state: for each node, return dp[sign][last_inv_dist]
    // sign: 0 = +1, 1 = -1 (current sign including all ancestor inversions)
    // dp[v][flip][d] = max additional sum from v's subtree given:
    //   flip = whether v's value is currently flipped (0=no, 1=yes)  
    //   d = distance to nearest ancestor that was inverted (capped at k)
    //
    // But this might be large. Instead, use iterative DFS.
    // 
    // State per node: dp[dist_to_last_inv][current_sign] = max sum of subtree
    // dist ranges from 0 to k (k+1 values), sign is binary
    // We can choose to invert node v or not (subject to dist constraint).
    
    // Iterative post-order DFS
    let mut parent = vec![usize::MAX; n];
    let mut order = Vec::with_capacity(n);
    let mut stack = vec![0usize];
    let mut visited = vec![false; n];
    visited[0] = true;
    while let Some(u) = stack.pop() {
      order.push(u);
      for &v in &adj[u] {
        if !visited[v] {
          visited[v] = true;
          parent[v] = u;
          stack.push(v);
        }
      }
    }
    
    // dp[v] = array of size (k+1): dp[v][d] = max subtree sum when the nearest ancestor inversion
    // is at distance d from v (d=k means "no constraint" i.e. >=k away or none)
    // The returned dp[v][d] accounts for optionally inverting v (if d >= k).
    let inf = i64::MIN / 2;
    let _dp: Vec<Vec<[i64; 2]>> = vec![vec![[0, 0]; k + 1]; n];
    // dp[v][d][s] = max sum for subtree of v, given nearest ancestor-inversion is d steps above,
    //   and s=0 means sign=+1, s=1 means sign=-1 (cumulative from ancestor inversions)
    // Actually simplify: dp[v][d] = max subtree sum for v's subtree given nearest ancestor
    //   inversion distance = d (0 = just inverted at v's parent, k = far/none)
    // We choose: invert v (need d >= k) or not invert.
    // After choosing, children get distance = (if we inverted v then 1 else d+1, capped at k).
    
    // Let me redefine: process in reverse order (post-order)
    // dp[v][d] = max subtree value sum when distance to last inversion above v = d
    //   (d=0 means parent was just inverted, d=k means >=k or none)
    
    // sign(v) = product of (number of inversions on root->v path) % 2
    // When we "invert" at v, all descendant values get multiplied by -1 additionally.
    
    // For the DFS state, track:
    //   dp[v][d] where d = min(distance_to_nearest_ancestor_inversion, k)
    //   This has (k+1) states.
    
    // At leaf v with sign s (0=+, 1=-):
    //   no invert: sum = nums[v] * (1 if s==0 else -1); next_d doesn't matter (leaf)
    //   invert (allowed if d >= k): sum = -nums[v] * (1 if s==0 else -1)
    //   dp[v][d] = max(option_no_inv, option_inv if d>=k)
    //   But we also need to propagate sign to children...
    
    // This requires tracking sign as well. Use dp[v][d][s]:
    //   s=0: current cumulative sign is +1
    //   s=1: current cumulative sign is -1
    // Size: n * (k+1) * 2
    
    let kp1 = k + 1;
    let mut memo: Vec<Vec<[i64; 2]>> = vec![vec![[inf, inf]; kp1]; n];
    
    // Process in reverse BFS order (leaves first)
    for &v in order.iter().rev() {
      let children: Vec<usize> = adj[v].iter().copied().filter(|&u| u != parent[v]).collect();
      
      for d in 0..kp1 {
        for s in 0..2usize {
          let base_val = if s == 0 { nums[v] as i64 } else { -(nums[v] as i64) };
          
          // Option 1: don't invert v
          // Children propagate with d+1 (capped at k) and same sign s
          let child_d = (d + 1).min(k);
          let no_inv = base_val + children.iter().map(|&c| memo[c][child_d][s]).sum::<i64>();
          
          // Option 2: invert v (only if d >= k)
          let inv = if d >= k {
            // After inverting v, sign flips, children get d=1, next ancestor inv = v
            let new_s = 1 - s;
            -base_val + children.iter().map(|&c| memo[c][1.min(k)][new_s]).sum::<i64>()
          } else {
            inf
          };
          
          memo[v][d][s] = no_inv.max(inv);
        }
      }
    }
    
    // Answer: start at root 0, d=k (no ancestor inversion), sign=0
    memo[0][k][0]
  }
}