3004. Maximum Subtree of the Same Color

MediumTreeDepth-First SearchArrayDynamic Programming

Problem Description

In this problem, we're presented with a tree made up of n nodes, and each node has a unique integer assigned to it as its color. The tree is defined by a list of edges, where each edge connects two nodes, forming a parent-child relationship. The objective is to find the largest possible subtree where all nodes share the same color.

A subtree refers to a node and all of its descendants in the tree. The size of a subtree is the total number of nodes it contains. Our goal is to determine the maximum size of any subtree such that all the nodes in that subtree have the exact same color.

To solve the problem, we must find a node v with the following property: Every node in the subtree with v as its root has the same color. The answer we're looking for is the maximum number of nodes in any such subtree.

The tree is rooted at node 0, and given to us as a 2D integer array called edges, where edges[i] = [u_i, v_i] indicates an edge connecting nodes u_i and v_i. We're also provided with a 0-indexed integer array colors where colors[i] gives us the color assigned to node i.

Intuition

Our approach to finding the maximum subtree with the same color utilizes Depth-First Search (DFS). The reasoning behind this is straightforward: To assess whether a particular node's subtree satisfies the color uniformity condition, we need to look at all of its descendants. A DFS allows us to explore each branch of the tree fully before moving to a different branch.

We start by constructing a graph g using an adjacency list, where each index represents a node, and the list at that index contains all of the node's children. Additionally, we maintain a size array, where size[a] is the size of the subtree rooted at node a. This will help us determine the size of each valid subtree.

To check for color uniformity, we define a dfs function that explores the tree and returns a boolean indicating whether all nodes in the current subtree (rooted at a) have the same color. While performing DFS, we traverse through each node's adjacent nodes (its children), recursively calling dfs on them, if they are not the parent fa.

As we perform the DFS, we carry two checks:

Whether the current node a has the same color as its adjacent node b.
Whether the subtrees rooted at b are uniform in color.

Only if both checks are true, does it mean that the current subtree rooted at a meets our criteria of uniform color. Each time we find such a subtree, we compare its size to our current maximum ans and update ans if we find a larger uniform-color subtree.

Finally, we begin our search from the root (node 0), initiating the DFS. The value of ans, which is continuously updated during DFS whenever a larger subtree fulfilling the requirement is found, will hold the size of the largest same-color subtree upon completion of the algorithm's execution.

Learn more about Tree, Depth-First Search and Dynamic Programming patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece：

Which of the following problems can be solved with backtracking (select multiple)

Generating subsetsGenerating random numbersGenerating permutationsSorting integers

Solution Approach

The solution presented here takes advantage of both a Depth-First Search (DFS) algorithm and recursion to efficiently traverse the tree and identify the maximum subtree with the same color.

The initial step involves creating an adjacency list g to represent the tree, leveraging the given edges to ascertain a node's children. This data structure is chosen for its ease of traversing connected nodes in a graph-like structure. Consequently, from each node (represented by an index in the list), we can quickly access all its connected nodes. In this context, a node a's children are represented by the list at g[a].

Next, we construct an array size with an initial value of 1 for every node a. As we traverse the tree and discover subtrees with matching colors, we'll cumulatively add the sizes of each subtree rooted at node a's children to size[a]. The size array is pivotal as it enables us to keep track of subtree sizes efficiently without recalculating them.

To implement the DFS, we define a recursive function dfs(a, fa) that performs the following steps:

Initialize a boolean variable ok to true. This variable signifies whether all nodes in the subtree rooted at node a have the same color.
Iterate over all adjacent nodes b of node a. If b is not the parent fa of a (to avoid cycling back up the tree), we recursively call dfs(b, a). The result is stored in t.
After each recursive call, we update ok by performing a logical AND operation with the current value of ok, the condition colors[a] == colors[b], and the return value t. This update reflects whether the same color is maintained throughout the current subtree and its descendants.
Additionally, during this traversal, we update size[a] by adding size[b] to accumulate the total size of all subtrees rooted at a.
If the result ok is true, indicating that the subtree rooted at a fulfills our condition (all nodes have the same color), we then update the global variable ans. ans represents the maximum size of any uniform-color subtree discovered so far. We check if size[a] is greater than the current ans and update ans accordingly.
Lastly, the function dfs returns the value of ok to indicate to upper levels of the recursive call whether the current subtree retains the same color.

Finally, we start the DFS with the root node 0 and an invalid parent node -1. After completing the DFS, ans, which has been updated during the process with the largest same-color subtree's size, is returned as the final result.

Through this approach, the solution elegantly combines an understanding of tree traversal, recursion, and the management of global and local state (with size and ok) to solve the problem efficiently.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Given a sorted array of integers and an integer called target, find the element that equals to the target and return its index. Select the correct code that fills the ___ in the given code snippet.

1def binary_search(arr, target):
2    left, right = 0, len(arr) - 1
3    while left ___ right:
4        mid = (left + right) // 2
5        if arr[mid] == target:
6            return mid
7        if arr[mid] < target:
8            ___ = mid + 1
9        else:
10            ___ = mid - 1
11    return -1
12

1public static int binarySearch(int[] arr, int target) {
2    int left = 0;
3    int right = arr.length - 1;
4
5    while (left ___ right) {
6        int mid = left + (right - left) / 2;
7        if (arr[mid] == target) return mid;
8        if (arr[mid] < target) {
9            ___ = mid + 1;
10        } else {
11            ___ = mid - 1;
12        }
13    }
14    return -1;
15}
16

1function binarySearch(arr, target) {
2    let left = 0;
3    let right = arr.length - 1;
4
5    while (left ___ right) {
6        let mid = left + Math.trunc((right - left) / 2);
7        if (arr[mid] == target) return mid;
8        if (arr[mid] < target) {
9            ___ = mid + 1;
10        } else {
11            ___ = mid - 1;
12        }
13    }
14    return -1;
15}
16

<=, left, right<, left, right<=, right, left<, right, left

I don’t know

Example Walkthrough

Assuming we have a tree with n = 5 nodes and the following edges and colors:

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

We begin by constructing our adjacency list g and size array from the information above:

1g = [[1, 2], [0, 3, 4], [0], [1], [1]]
2size = [1, 1, 1, 1, 1] // Initialize all sizes to 1

Now let's perform the Depth-First Search (DFS):

We start DFS from the root node 0 with an invalid parent -1. The children of node 0 are [1, 2]. Since colors[0] is 1, we will compare with the colors of its children.
- For child 1, the color matches the root. We recursively call DFS on node 1 with parent 0. Here, ok starts as true.
- Within this call, node 1 has two children: 3 and 4. Both share the same color as 1. Thus, we add their sizes to size[1] and ok remains true for both.
- After traversing children 3 and 4, size[1] becomes 1 + 1 + 1 = 3 (node 1 + 3 and 4) and the function returns true.
- For child 2, the color does not match the root (colors[2] is 2). We call DFS on node 2 with parent 0. Here, ok is false immediately because the color is different. There are no further children to explore, and size[2] remains 1.
Now back at node 0, we have the results from its children. Since the child 1 returned true and has matching colors, the subtree size rooted at 0 that has a uniform color is size[1] = 3. However, since the child 2 had a different color, ok for node 0 is false.
Our maximum answer ans updates after visiting node 1’s subtree to the maximum subtree size found, which is 3 in this case. Node 0 does not meet our criteria because it has a different color child (node 2), so our final answer remains 3.

In this walkthrough, we navigated the tree in a depth-first manner, ensuring we fully understood each node's subtree before moving on. We successfully found that the largest subtree where all nodes shared the same color was rooted at node 1, with a size of 3.

Solution Implementation

1from typing import List
2
3class Solution:
4    def maximumSubtreeSize(self, edges: List[List[int]], colors: List[int]) -> int:
5        # A depth-first search function that traverses the graph and
6        # computes the size of each subtree with a single color.
7        def dfs(node: int, parent: int) -> bool:
8            is_uniform_color = True  # A flag to check if the subtree has a uniform color
9            # Iterate over all neighboring nodes
10            for neighbor in graph[node]:
11                # If neighbor is not the parent, then it's part of the subtree
12                if neighbor != parent:
13                    # Perform a DFS on the child node
14                    is_subtree_uniform = dfs(neighbor, node)
15                    # The current subtree can only be uniform if its children are uniform
16                    # and the colors match
17                    is_uniform_color = is_uniform_color and colors[node] == colors[neighbor] and is_subtree_uniform
18                    # Add the size of the child's subtree to the current node's subtree
19                    subtree_size[node] += subtree_size[neighbor]
20            # If the current node's subtree is uniform, check if it's the maximum seen so far
21            if is_uniform_color:
22                nonlocal max_subtree_size
23                max_subtree_size = max(max_subtree_size, subtree_size[node])
24            return is_uniform_color
25
26        # Initialize the number of nodes in the graph
27        num_nodes = len(edges) + 1
28        # Create an adjacency list for the graph
29        graph = [[] for _ in range(num_nodes)]
30        # Initialize the subtree size list with all ones (each node is a subtree of size 1)
31        subtree_size = [1] * num_nodes
32        # Build the graph connections from the given edges
33        for a, b in edges:
34            graph[a].append(b)
35            graph[b].append(a)
36        # Initialize the answer to zero
37        max_subtree_size = 0
38        # Start the DFS from the first node (assuming 0-indexed) with no parent (-1)
39        dfs(0, -1)
40        # Return the maximum size of a subtree with uniform color
41        return max_subtree_size
42

1class Solution {
2    private List<Integer>[] adjList; // Adjacency list for representing the graph.
3    private int[] nodeColors;       // Array to store colors of the nodes.
4    private int[] subtreeSize;      // Array to store sizes of the subtrees.
5    private int maxSubtreeSize;     // Variable to keep track of the maximum subtree size found.
6
7    // Method to calculate the maximum subtree size where all nodes have the same color.
8    public int maximumSubtreeSize(int[][] edges, int[] colors) {
9        int n = edges.length + 1; // Total number of nodes.
10        adjList = new List[n];
11        subtreeSize = new int[n];
12        nodeColors = colors;
13        Arrays.fill(subtreeSize, 1); // Initialize all subtree sizes to 1 (each node at least has a size of 1 - itself).
14        Arrays.setAll(adjList, i -> new ArrayList<>());// Initialize lists to represent adjacency.
15
16        // Build the graph from the edge list.
17        for (int[] edge : edges) {
18            int from = edge[0], to = edge[1];
19            adjList[from].add(to);
20            adjList[to].add(from);
21        }
22
23        // Perform Depth-First Search starting from node 0.
24        dfs(0, -1);
25        return maxSubtreeSize;
26    }
27
28    // Depth-First Search method to explore the graph and calculate subtree sizes.
29    private boolean dfs(int node, int parent) {
30        boolean isMonochrome = true; // Flag to check if the current subtree contains the same color.
31        // Iterate over all the neighbors of the current node.
32        for (int neighbor : adjList[node]) {
33            // If neighbor is not the parent.
34            if (neighbor != parent) {
35                // Perform DFS on the neighboring node.
36                boolean isNeighborMonochrome = dfs(neighbor, node);
37                // Check if the current node and its neighbor have the same color and the neighbor is monochrome.
38                isMonochrome = isMonochrome && nodeColors[node] == nodeColors[neighbor] && isNeighborMonochrome;
39                // Update the size of the current subtree by adding the size of the neighboring subtree.
40                subtreeSize[node] += subtreeSize[neighbor];
41            }
42        }
43        // If the current subtree is monochrome, update the maximum subtree size found so far.
44        if (isMonochrome) {
45            maxSubtreeSize = Math.max(maxSubtreeSize, subtreeSize[node]);
46        }
47        // Return whether the current subtree is monochrome.
48        return isMonochrome;
49    }
50}
51

1#include <vector>
2#include <functional>
3using namespace std;
4
5class Solution {
6public:
7    // Function returning the size of the maximum subtree with uniform colors.
8    int maximumSubtreeSize(vector<vector<int>>& edges, vector<int>& colors) {
9        int numOfNodes = edges.size() + 1;  // Calculate the number of nodes.
10        vector<vector<int>> graph(numOfNodes); // Adjacency list representation of the graph.
11        vector<int> subtreeSize(numOfNodes, 1);  // Initialize all subtree sizes to 1 (each node).
12
13        // Building the undirected graph from the edges.
14        for (const auto& edge : edges) {
15            int nodeA = edge[0], nodeB = edge[1];
16            graph[nodeA].push_back(nodeB);
17            graph[nodeB].push_back(nodeA);
18        }
19
20        int maxSubtreeSize = 0;  // This will hold the result - the size of the largest subtree.
21      
22        // Recursive DFS function to traverse the graph while calculating subtree sizes.
23        // It returns true if all nodes in the current subtree have the same color.
24        function<bool(int, int)> depthFirstSearch = [&](int node, int parent) {
25            bool isUniformColor = true;  // To check if all children have the same color as the current node.
26
27            // Traverse all neighbors of the current node.
28            for (int neighbor : graph[node]) {
29                // If neighbor is not the parent, do DFS on the neighbor.
30                if (neighbor != parent) {
31                    bool subtreeIsUniformColor = depthFirstSearch(neighbor, node);
32                    // Check if the neighbor's subtree is uniformly colored and has the same color as the current node.
33                    isUniformColor = isUniformColor && colors[node] == colors[neighbor] && subtreeIsUniformColor;
34                    subtreeSize[node] += subtreeSize[neighbor];  // Update the size of the current subtree.
35                }
36            }
37
38            // If the current subtree is uniformly colored, update the maximum subtree size.
39            if (isUniformColor) {
40                maxSubtreeSize = max(maxSubtreeSize, subtreeSize[node]);
41            }
42            return isUniformColor;
43        };
44
45        // Start DFS from the root node (assuming it is labeled with 0) with no parent (-1).
46        depthFirstSearch(0, -1);
47
48        // Return the size of the largest uniformly colored subtree found.
49        return maxSubtreeSize;
50    }
51};
52

1function maximumSubtreeSize(edges: number[][], colors: number[]): number {
2    // The number of nodes in the tree.
3    const numberOfNodes = edges.length + 1;
4    // The adjacency list to represent the tree graph.
5    const graph: number[][] = Array.from({ length: numberOfNodes }, () => []);
6
7    // Fill the adjacency list with the edges.
8    for (const [node1, node2] of edges) {
9        graph[node1].push(node2);
10        graph[node2].push(node1);
11    }
12
13    // Array to store the size of each subtree.
14    const subtreeSizes: number[] = Array(numberOfNodes).fill(1);
15    // Variable to keep track of the size of the maximum monochromatic subtree.
16    let maxSubtreeSize = 0;
17
18    // Recursive depth-first search function to traverse graph and calculate subtree sizes.
19    const depthFirstSearch = (currentNode: number, parent: number): boolean => {
20        // Flag to check if current subtree is monochromatic.
21        let isMonochromatic = true;
22
23        // Traverse all adjacent nodes.
24        for (const adjacentNode of graph[currentNode]) {
25            // If adjacent node is not the parent.
26            if (adjacentNode !== parent) {
27                // Recurse deeper into the tree.
28                const isChildMonochromatic = depthFirstSearch(adjacentNode, currentNode);
29                // Update the monochromatic status of the current node.
30                isMonochromatic = isMonochromatic && isChildMonochromatic && colors[currentNode] === colors[adjacentNode];
31                // Aggregate the size of the subtree.
32                subtreeSizes[currentNode] += subtreeSizes[adjacentNode];
33            }
34        }
35
36        // If subtree rooted at current node is monochromatic, update the maximum size if needed.
37        if (isMonochromatic) {
38            maxSubtreeSize = Math.max(maxSubtreeSize, subtreeSizes[currentNode]);
39        }
40
41        // Return the monochromatic status of the subtree rooted at current node.
42        return isMonochromatic;
43    };
44
45    // Start the depth-first search from node 0 with parent -1 (as there is no parent for root).
46    depthFirstSearch(0, -1);
47
48    // Return the size of the largest monochromatic subtree.
49    return maxSubtreeSize;
50}
51

Not Sure What to Study? Take the 2-min Quiz：

What's the relationship between a tree and a graph?

A graph is a special treeA tree is a special graphNo relationshipThey are the same thing

I don’t know

Time and Space Complexity

Time Complexity

The time complexity of the code is O(n). This is because the function dfs is called recursively for each node in the tree exactly once. During each call, the function processes the current node, and the processing time for each node is proportional to the number of its direct children due to the loop for b in g[a]. Since the graph is a tree represented by n - 1 edges, the total number of such direct connections is n - 1. Thus, the overall time to process all nodes once is proportional to n, hence the time complexity is linear with respect to the number of nodes.

Space Complexity

The space complexity of the code is also O(n). The g list of lists (which represents the adjacency list of the tree) and the size array each consume space proportional to n. Additionally, the recursion stack for the depth-first search may also grow up to O(n) in the case of a degenerate (linked-list-like) tree where each node has only one child except for the leaf node. Thus, the space used by the data structures combined with the recursion stack's depth accounts for the total space complexity of O(n).

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

What is the running time of the following code?

1int sqrt(int n) {
2  for (int guess = 1; guess * guess <= n; guess++) {
3    if (guess * guess == n) {
4      return guess;
5    }
6  }
7  return -1;
8}

O(1)O(log n)O(sqrt(n))O(n)

I don’t know

Overview

Basic Data Structures

Basic Algorithms

Overview

Sorted Array

Implicitly Sorted Array

Advanced

Speedrun

Introduction

DFS on Tree

Binary Search Tree

Advanced

Combinatorial Search

Pruning

Additional States

Aggregation and Memoization

Dedup

Additional Practices

Speedrun

Introduction

BFS on Tree

Introduction

Vanilla BFS

Matrix as Graph

Implicit Graph

Directed Graph / Topological Sort

Weighted Graph

Minimum Spanning Tree

Speedrun

Introduction

Same Direction

Opposite Direction

Sliding Window

Prefix Sum

Cycle Finding

Advanced

Speedrun

Introduction

Top K

Moving Best

Multiple Heaps

Introduction

Grid

Dynamic number of subproblems

Interval

Dual-Sequence

Knapsack

Bitmask

Additional Practices

Disjoint Set Union | Union Find

Trie

Data Structure Design

Segment Tree

Interval

Monotonic Stack

Divide and Conquer

Line Sweep

Tree Traversal without Recursion

Greedy

Math

Matrix

Classics

Concepts

Projects

Amazon OA

Microsoft OA

Google OA

Twitter OA

Editorials

3004. Maximum Subtree of the Same Color

Problem Description

Intuition

Solution Approach

Example Walkthrough

Solution Implementation

Time and Space Complexity

Time Complexity

Space Complexity

Recommended Readings