Skip to main content
Back to problems
#3575
Hard Algorithms

Maximum good subtree score

Array Dynamic Programming Bit Manipulation Tree Depth-First Search Bitmask
45.5% acceptance
Feb 25, 2026
51
8
Given a tree rooted at 0, n nodes, each with value vals[i]. A subset of a subtree is "good" if every digit 0-9 appears at most once across all selected node values. Score = sum of values. maxScore[u] = max score of good subset in subtree of u. Return sum of all maxScore[u] mod 10^9+7.

Solution

Rust
Time O(n * m)
Space O(n * m)
LeetCode
solution.rs
impl Solution {
  pub fn good_subtree_sum(vals: Vec<i32>, par: Vec<i32>) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let n = vals.len();

    // Precompute digit mask for each node (which digits appear in vals[i])
    // If a digit appears more than once in vals[i], the node cannot be part of any good subset.
    let digit_mask: Vec<Option<u16>> = vals
      .iter()
      .map(|&v| {
        let mut mask = 0u16;
        let mut valid = true;
        let mut x = v;
        while x > 0 {
          let d = (x % 10) as u16;
          if mask & (1 << d) != 0 {
            valid = false;
            break;
          }
          mask |= 1 << d;
          x /= 10;
        }
        if valid { Some(mask) } else { None }
      })
      .collect();

    // Build children list
    let mut children: Vec<Vec<usize>> = vec![vec![]; n];
    for i in 1..n {
      let p = par[i] as usize;
      children[p].push(i);
    }

    // DFS post-order. For each subtree, compute:
    // dp[mask] = max score of a good subset with these digits used,
    //            considering only nodes in this subtree.
    // When merging child into parent:
    //   for each (child_mask, child_score) and (cur_mask, cur_score):
    //     if (child_mask & cur_mask) == 0: new state (child_mask|cur_mask, child_score+cur_score)
    // Also, node u itself can be included if digit_mask[u] doesn't conflict.

    let mut total_sum = 0i64;

    // Iterative post-order DFS
    let mut order = vec![];
    let mut stack = vec![0usize];
    while let Some(v) = stack.pop() {
      order.push(v);
      for &c in &children[v] {
        stack.push(c);
      }
    }
    order.reverse(); // now leaves first

    // dp[v] = HashMap or array of size 2^10=1024 of max score for each digit mask used
    let mut dp: Vec<Vec<i64>> = vec![vec![-1i64; 1024]; n];

    for &v in &order {
      let vm = digit_mask[v];
      let _ = vm; // used below
      // Start with empty subset (mask=0, score=0)
      dp[v][0] = 0;

      // Merge each child
      for &c in &children[v] {
        // Merge dp[v] with dp[c]
        let mut new_dp = vec![-1i64; 1024];
        for cur_mask in 0usize..1024 {
          if dp[v][cur_mask] < 0 { continue; }
          // Option 1: don't take anything from child
          if new_dp[cur_mask] < dp[v][cur_mask] {
            new_dp[cur_mask] = dp[v][cur_mask];
          }
          // Option 2: take some subset from child
          let avail = (!cur_mask) & 0x3FF;
          let mut child_mask = avail;
          loop {
            if dp[c][child_mask] >= 0 {
              let combined = cur_mask | child_mask;
              let score = dp[v][cur_mask] + dp[c][child_mask];
              if new_dp[combined] < score {
                new_dp[combined] = score;
              }
            }
            if child_mask == 0 { break; }
            child_mask = (child_mask - 1) & avail;
          }
        }
        dp[v] = new_dp;
      }

      // Now try including node v itself (only if its digit mask is valid and no conflict)
      if let Some(vm) = digit_mask[v] {
        let vm = vm as usize;
        let mut new_dp = dp[v].clone();
        for cur_mask in 0usize..1024 {
          if dp[v][cur_mask] < 0 { continue; }
          if cur_mask & vm == 0 {
            let combined = cur_mask | vm;
            let score = dp[v][cur_mask] + vals[v] as i64;
            if new_dp[combined] < score {
              new_dp[combined] = score;
            }
          }
        }
        dp[v] = new_dp;
      }

      // maxScore[v] = max over all dp[v][mask]
      let max_score = *dp[v].iter().max().unwrap();
      total_sum = (total_sum + max_score % MOD) % MOD;
    }

    total_sum as i32
  }
}