LeetCode 3004 - Maximum Subtree of the Same Color

We are given a rooted tree with n nodes, where node 0 is the root. The input edges describes the tree structure. Each entry edges[i] = [u, v] indicates that there is an undirected edge between nodes u and v.

leetcodemediumarraydynamic-programmingtreedepth-first-search

LeetCode Problem 3004

Difficulty: 🟡 Medium
Topics: Array, Dynamic Programming, Tree, Depth-First Search

Solution

Problem Understanding

We are given a rooted tree with n nodes, where node 0 is the root.

The input edges describes the tree structure. Each entry edges[i] = [u, v] indicates that there is an undirected edge between nodes u and v. Since the problem guarantees that the graph is a tree rooted at node 0, every node belongs to exactly one subtree rooted somewhere in the tree.

We are also given an array colors, where colors[i] represents the color assigned to node i.

The goal is to find a node v such that every node inside the subtree rooted at v has exactly the same color. Among all such valid subtrees, we must return the size of the largest one.

A subtree rooted at node v contains:

  • Node v itself.
  • Every descendant of v.

For a subtree to be valid, every node inside it must share one identical color.

The constraints are important:

  • n can be as large as 50,000.
  • A subtree can also contain up to 50,000 nodes.
  • Colors can be large values (<= 100000), so color values themselves cannot be used as array indices in a compact DP table.
  • Since the tree is large, repeatedly traversing subtrees would be too expensive.

These constraints strongly suggest that we need a single traversal of the tree, most naturally a Depth-First Search (DFS).

Some important edge cases include:

  • A tree consisting of only one node. The answer must be 1.
  • A tree where every node has the same color. The answer is the size of the entire tree.
  • A tree where every parent has children of different colors. In that case, only leaf nodes form valid same-color subtrees, so the answer is 1.
  • Deep chains of nodes where recursion depth could become large in Python.
  • Internal nodes whose descendants are all the same color but differ from the root of that subtree. Such subtrees are not valid because every node, including the root, must share the same color.

Approaches

Brute Force

A straightforward approach would be to examine every node as a potential subtree root.

For each node:

  1. Traverse its entire subtree.
  2. Collect all colors appearing in that subtree.
  3. Check whether all colors are identical.
  4. If yes, update the answer with the subtree size.

This approach is correct because it explicitly verifies the condition for every subtree.

However, it is extremely inefficient. Consider a chain of n nodes. The subtree of the root contains n nodes, the next subtree contains n-1, and so on.

The total work becomes:

$$n + (n-1) + (n-2) + \cdots + 1 = O(n^2)$$

With n = 50,000, this is far too slow.

Key Insight

A subtree is monochromatic if and only if:

  1. Every child subtree is monochromatic.
  2. Every child subtree has the same color as the current node.

This observation allows us to process the tree bottom-up.

During a DFS, each node can return:

  • Whether its subtree is monochromatic.
  • What color that monochromatic subtree has.
  • The size of the subtree.

Leaf nodes are automatically monochromatic.

For an internal node:

  • If any child subtree is not monochromatic, the current subtree cannot be monochromatic.
  • If a child subtree has a different color from the current node, the current subtree cannot be monochromatic.
  • Otherwise, the current subtree is monochromatic.

Since each edge is processed exactly once, we obtain a linear-time solution.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n²) O(n) Repeatedly traverses many overlapping subtrees
Optimal DFS O(n) O(n) Processes each node once using postorder traversal

Algorithm Walkthrough

Optimal DFS Strategy

  1. Build an adjacency list from the edge list.

Since the input describes a tree, an adjacency list allows efficient traversal of neighboring nodes. 2. Start a DFS from the root node 0.

We perform a postorder traversal because a node's validity depends on the results of its children. 3. For each node, initialize:

  • subtree_size = 1
  • same_color = True

Every node contributes itself to its subtree. 4. Visit each child.

Skip the parent node to avoid traversing backward in the undirected tree. 5. Recursively process the child.

The child returns:

  • Whether its subtree is monochromatic.
  • Its subtree size.
  1. Add the child's size to the current subtree size.

This allows us to compute subtree sizes in one traversal. 7. Verify the monochromatic condition.

If the child subtree is not monochromatic, mark the current subtree as invalid.

If the child's color differs from the current node's color, mark the current subtree as invalid. 8. After processing all children:

  • If the subtree remains valid, update the global answer using its size.
  • Return that this subtree is monochromatic.
  • Otherwise return that it is not monochromatic.
  1. Continue until the DFS finishes at the root.
  2. Return the maximum valid subtree size encountered.

Why it works

The algorithm processes nodes in postorder, meaning every child subtree is completely analyzed before its parent. A subtree rooted at node u is monochromatic exactly when every child subtree is monochromatic and all those child subtrees share the same color as u. Because these conditions are both necessary and sufficient, every subtree is classified correctly. Since every node is visited once, the largest valid subtree is found during the traversal.

Python Solution

from typing import List
import sys

class Solution:
    def maximumSubtreeSize(self, edges: List[List[int]], colors: List[int]) -> int:
        n = len(colors)

        sys.setrecursionlimit(100000)

        graph = [[] for _ in range(n)]
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        answer = 1

        def dfs(node: int, parent: int):
            nonlocal answer

            subtree_size = 1
            is_monochromatic = True

            for neighbor in graph[node]:
                if neighbor == parent:
                    continue

                child_mono, child_size = dfs(neighbor, node)

                subtree_size += child_size

                if not child_mono:
                    is_monochromatic = False

                if colors[neighbor] != colors[node]:
                    is_monochromatic = False

            if is_monochromatic:
                answer = max(answer, subtree_size)

            return is_monochromatic, subtree_size

        dfs(0, -1)
        return answer

The implementation first constructs an adjacency list representation of the tree.

The DFS returns two pieces of information:

  • Whether the subtree rooted at the current node is monochromatic.
  • The size of that subtree.

Each node starts as a valid monochromatic subtree of size one. As children are processed, their subtree sizes are accumulated.

The critical observation is that if a child subtree is monochromatic and has the same color as the current node, then every node inside that child subtree must also have the same color as the current node. Therefore checking the child's root color is sufficient.

Whenever a node's subtree satisfies the monochromatic condition, the global answer is updated.

Because every edge is traversed exactly once, the algorithm runs in linear time.

Go Solution

func maximumSubtreeSize(edges [][]int, colors []int) int {
	n := len(colors)

	graph := make([][]int, n)
	for _, e := range edges {
		u := e[0]
		v := e[1]

		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	answer := 1

	var dfs func(int, int) (bool, int)

	dfs = func(node int, parent int) (bool, int) {
		subtreeSize := 1
		isMonochromatic := true

		for _, neighbor := range graph[node] {
			if neighbor == parent {
				continue
			}

			childMono, childSize := dfs(neighbor, node)

			subtreeSize += childSize

			if !childMono {
				isMonochromatic = false
			}

			if colors[neighbor] != colors[node] {
				isMonochromatic = false
			}
		}

		if isMonochromatic {
			if subtreeSize > answer {
				answer = subtreeSize
			}
		}

		return isMonochromatic, subtreeSize
	}

	dfs(0, -1)
	return answer
}

The Go solution mirrors the Python implementation almost exactly.

The adjacency list is represented using a slice of slices. The DFS closure returns both the monochromatic status and subtree size. Since the answer is maintained globally, there is no need to return it through recursion.

Integer overflow is not a concern because subtree sizes are at most 50,000, which easily fits into Go's int.

Worked Examples

Example 1

edges  = [[0,1],[0,2],[0,3]]
colors = [1,1,2,3]

Tree:

      0(1)
    /  |  \
 1(1) 2(2) 3(3)

DFS processing order:

Node Subtree Size Monochromatic
1 1 True
2 1 True
3 1 True

Current answer becomes 1.

Now process node 0.

Child Child Color Parent Color Match
1 1 1 Yes
2 2 1 No
3 3 1 No

Node 0 is not monochromatic.

Final answer:

1

Example 2

edges  = [[0,1],[0,2],[0,3]]
colors = [1,1,1,1]

Tree:

      0
    / | \
   1  2  3

Leaves:

Node Size Valid
1 1 True
2 1 True
3 1 True

Process node 0.

Child Child Color Parent Color
1 1 1
2 1 1
3 1 1

All conditions hold.

Subtree size:

1 + 1 + 1 + 1 = 4

Answer becomes:

4

Example 3

edges  = [[0,1],[0,2],[2,3],[2,4]]
colors = [1,2,3,3,3]

Tree:

        0(1)
       /   \
    1(2)   2(3)
          /   \
       3(3) 4(3)

Process leaves first.

Node Size Valid
1 1 True
3 1 True
4 1 True

Process node 2.

Child Color
3 3
4 3

Both match node 2's color.

Subtree size:

1 + 1 + 1 = 3

Answer becomes 3.

Process node 0.

Child colors:

2 and 3

Neither matches color 1.

Node 0 is invalid.

Final answer:

3

Complexity Analysis

Measure Complexity Explanation
Time O(n) Every node and edge is processed once
Space O(n) Adjacency list and recursion stack

The DFS performs a single traversal of the tree. Since a tree with n nodes contains exactly n - 1 edges, every edge is visited only once in each direction. The adjacency list requires O(n) space, and in the worst case a skewed tree can produce a recursion depth of O(n).

Test Cases

from typing import List

s = Solution()

assert s.maximumSubtreeSize(
    [[0,1],[0,2],[0,3]],
    [1,1,2,3]
) == 1  # provided example 1

assert s.maximumSubtreeSize(
    [[0,1],[0,2],[0,3]],
    [1,1,1,1]
) == 4  # provided example 2

assert s.maximumSubtreeSize(
    [[0,1],[0,2],[2,3],[2,4]],
    [1,2,3,3,3]
) == 3  # provided example 3

assert s.maximumSubtreeSize(
    [],
    [7]
) == 1  # single node tree

assert s.maximumSubtreeSize(
    [[0,1]],
    [1,2]
) == 1  # two different colors

assert s.maximumSubtreeSize(
    [[0,1]],
    [5,5]
) == 2  # entire tree valid

assert s.maximumSubtreeSize(
    [[0,1],[1,2],[2,3]],
    [1,1,1,1]
) == 4  # chain with same color

assert s.maximumSubtreeSize(
    [[0,1],[1,2],[2,3]],
    [1,2,3,4]
) == 1  # chain with all different colors

assert s.maximumSubtreeSize(
    [[0,1],[0,2],[1,3],[1,4]],
    [2,2,2,2,2]
) == 5  # entire tree monochromatic

assert s.maximumSubtreeSize(
    [[0,1],[0,2],[2,3],[2,4]],
    [1,1,2,2,2]
) == 3  # largest valid subtree is internal node

Test Summary

Test Why
Example 1 Multiple colors at root
Example 2 Entire tree has same color
Example 3 Largest valid subtree is internal
Single node Smallest possible input
Two nodes, different colors No subtree larger than 1
Two nodes, same color Whole tree valid
Uniform chain Deep monochromatic subtree
Different-color chain Every valid subtree is a leaf
Entire tree same color Large valid root subtree
Internal maximum subtree Ensures non-root answer is found

Edge Cases

Single Node Tree

When n = 1, there are no edges and the only subtree is the node itself. A common bug is assuming every node has children or failing to initialize the answer properly. The implementation initializes answer = 1, so this case is handled naturally.

Entire Tree Has One Color

If every node has the same color, the correct answer is the size of the entire tree. Some implementations only verify immediate children and forget that deeper descendants must also be checked. Our DFS guarantees that every descendant subtree is validated before its parent, so the root is correctly recognized as a monochromatic subtree.

Largest Valid Subtree Is Not the Root

Many trees contain a mixed-color root but a large monochromatic subtree deeper in the tree. An implementation that only checks the root would fail. The DFS evaluates every node independently and updates the global maximum whenever a valid subtree is found, ensuring internal subtrees are considered.

Deep Skewed Tree

A tree may degenerate into a linked list of length 50,000. Recursive DFS can exceed Python's default recursion limit. The implementation explicitly increases the recursion limit with sys.setrecursionlimit(100000), preventing recursion depth failures on valid inputs.

Mixed Valid and Invalid Child Subtrees

A node may have several children where some subtrees are monochromatic and others are not. A buggy implementation might stop after finding one valid child. Our DFS requires every child subtree to be monochromatic and every child root color to match the current node's color. If any child violates these conditions, the current subtree is marked invalid. This exactly matches the problem definition.