Maximum xor of two non overlapping subtrees

Tree Depth-First Search Graph Theory Trie

52.1% acceptance

Mar 31, 2026

There is an undirected tree with n nodes labeled from 0 to n - 1. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree. The root of the tree is the node labeled 0. Each node has an associated value. You are given an array values of length n, where values[i] is the value of the ith node. Select any two non-overlapping subtrees. Your score is the bitwise XOR of the sum of the values within those subtrees. Return the maximum possible score you can achieve. If it is impossible to find two nonoverlapping subtrees, return 0. Note that: The subtree of a node is the tree consisting of that node and all of its descendants. Two subtrees are non-overlapping if they do not share any common node.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn max_xor(n: i32, edges: Vec<Vec<i32>>, values: Vec<i32>) -> i64 {
    let n = n as usize;
    let mut adj = vec![vec![]; n];
    for e in &edges {
      let a = e[0] as usize;
      let b = e[1] as usize;
      adj[a].push(b);
      adj[b].push(a);
    }
    // Compute subtree sums
    let mut subtree_sum = vec![0i64; n];
    let mut order = Vec::with_capacity(n);
    let mut parent = vec![usize::MAX; 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);
        }
      }
    }
    for &u in order.iter().rev() {
      subtree_sum[u] += values[u] as i64;
      if parent[u] != usize::MAX {
        subtree_sum[parent[u]] += subtree_sum[u];
      }
    }
    
    // Trie for XOR maximization (47 bits)
    const BITS: usize = 47;
    let mut trie = vec![[0i32; 2]]; // trie[node] = [child0, child1], -1 means no child
    trie[0] = [-1, -1];
    
    let mut ans = 0i64;
    
    // DFS: on enter query trie, on exit insert into trie
    let mut children: Vec<Vec<usize>> = vec![vec![]; n];
    for u in 0..n {
      for &v in &adj[u] {
        if parent[v] == u {
          children[u].push(v);
        }
      }
    }
    
    // Iterative DFS
    let mut dfs_stack: Vec<(usize, usize)> = vec![(0, 0)]; // (node, next_child_index)
    // On first visit to node, query trie
    let mut entered = vec![false; n];
    
    while let Some(&mut (u, ref mut ci)) = dfs_stack.last_mut() {
      if !entered[u] {
        entered[u] = true;
        // Query trie if trie has entries
        if trie.len() > 1 || trie[0][0] != -1 || trie[0][1] != -1 {
          let s = subtree_sum[u];
          let mut node = 0i32;
          let mut xor_val = 0i64;
          let mut valid = true;
          for bit in (0..BITS).rev() {
            let b = ((s >> bit) & 1) as usize;
            let want = 1 - b;
            if trie[node as usize][want] != -1 {
              xor_val |= 1i64 << bit;
              node = trie[node as usize][want];
            } else if trie[node as usize][b] != -1 {
              node = trie[node as usize][b];
            } else {
              valid = false;
              break;
            }
          }
          if valid {
            ans = ans.max(xor_val);
          }
        }
      }
      if *ci < children[u].len() {
        let child = children[u][*ci];
        *ci += 1;
        dfs_stack.push((child, 0));
      } else {
        dfs_stack.pop();
        // Insert subtree_sum[u] into trie
        let s = subtree_sum[u];
        let mut node = 0usize;
        for bit in (0..BITS).rev() {
          let b = ((s >> bit) & 1) as usize;
          if trie[node][b] == -1 {
            trie.push([-1, -1]);
            trie[node][b] = (trie.len() - 1) as i32;
          }
          node = trie[node][b] as usize;
        }
      }
    }
    
    ans
  }
}