Find Maximum Subtree of the Same Color

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] = [u_i, vi] means there is an edge between the nodes v_i and u_i. Another array colors[] of size n is also given, where colors[i] is the color assigned to node i. We have to find a node v such that every node in the subtree of v has the same color. The task is to return the size of such subtree with the maximum number of nodes possible.

Example:

Input: edges = {{0,1},{0,2},{0,3}}, colors = {1,1,2,3}
Output: 1
Explanation: Each color is represented as: 1 -> Red, 2 -> Green, 3 -> Blue. We can see that the subtree rooted at node 0 has children with different colors. Any other subtree is of the same color and has a size of 1. Hence, we return 1.

Input: edges = {{0,1},{0,2},{0,3}}, colors = {1,1,1,1}
Output: 4
Explanation: The whole tree has the same color, and the subtree rooted at node 0 has the most number of nodes which is 4. Hence, we return 4.

Approach:

For every subtree inform the parent if all nodes in the subtree has same color. Carry three pieces of information in a post order traversal back up.

• Does the subtree below has all same colors
• Sum of all nodes in the subtree
• Color of the root node of subtree.

After every post order traversal on each subtree decide what will be the values of those three that you need to send up to the parent.

Steps-by-step approach:

• Initialize variable m_maxSize to 1, representing the maximum size of a subtree with all nodes of the same color.
• Define struct Res to represent the result of DFS, containing fields all_same_color, num_nodes, and color.
• Create unction dfs to perform DFS on a subtree:
  • If the node is a leaf (has no children), return result indicating all nodes have the same color, node count is 1, and color is the node's color.
  • Iterate over children of the node, recursively call dfs.
  • Update color_match indicating if all nodes in subtree have same color.
  • Update sum_of_nodes with sum of nodes in subtree.
  • If all nodes in subtree have same color, update m_maxSize.
  • Return result of DFS on subtree.
• Create function maximumSubtreeSize to compute maximum size of subtree with all nodes of same color:
  • Create graph as adjacency list from edges.
  • Call DFS on entire graph.
  • Return m_maxSize

Below is the implementation of the above approach:

2

Time complexity: O(n)
Auxiliary Space: O(n) considering the recursive stack space