Maximum subtree of the same color

Array Dynamic Programming Tree Depth-First Search

58.7% acceptance

Mar 31, 2026

You are given a 2D integer array edges representing a tree with n nodes, numbered from 0 to n - 1, rooted at node 0, where edges[i] = [ui, vi] means there is an edge between the nodes vi and ui. You are also given a 0-indexed integer array colors of size n, where colors[i] is the color assigned to node i. We want to find a node v such that every node in the subtree of v has the same color. Return the size of such subtree with the maximum number of nodes possible.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn maximum_subtree_size(edges: Vec<Vec<i32>>, colors: Vec<i32>) -> i32 {
    let n = colors.len();
    let mut adj = vec![vec![]; n];
    for e in &edges {
      adj[e[0] as usize].push(e[1] as usize);
      adj[e[1] as usize].push(e[0] as usize);
    }
    let mut ans = 1;
    // Iterative DFS with post-order processing
    // size[i] = subtree size, 0 means not uniform
    let mut size = vec![0i32; 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);
        }
      }
    }
    // Process in reverse order (post-order)
    for &u in order.iter().rev() {
      size[u] = 1; // leaf or start
      let mut uniform = true;
      for &v in &adj[u] {
        if v != parent[u] {
          if size[v] == 0 || colors[v] != colors[u] {
            uniform = false;
          } else {
            size[u] += size[v];
          }
        }
      }
      if !uniform {
        size[u] = 0;
      } else {
        ans = ans.max(size[u]);
      }
    }
    ans
  }
}