3004. Maximum Subtree of the Same Color

MediumTreeDepth-First SearchArrayDynamic Programming
Leetcode Link

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.

Flowchart Walkthrough

Here's how to use the algorithm flowchart (the Flowchart) to determine the appropriate algorithm for Leetcode problem 3004, Maximum Subtree of the Same Color:

  1. Is it a graph?

    • Yes: The tree can be represented as a special case of a graph where there are no cycles and each node (except the root) has exactly one parent.

  2. Is it a tree?

    • Yes: By its description, the problem deals with a tree since we are looking at subtrees which implies a hierarchical structure.

  3. DFS (Depth-First Search)

    • At this point, since we've identified the structure as a tree and the problem revolves around checking properties (same colors) within subtrees, DFS is a suitable algorithm. DFS allows exploring each branch of the tree completely before moving on to another branch, which is useful for problems involving subtrees or dependent substructures.

Conclusion: Leveraging DFS is beneficial for problems involving subtrees because it enables deep exploration, and can efficiently check or compute conditions or gather values as it recurses back up the tree. In this scenario, DFS can be used to traverse each node and maintain state or check conditions (like color continuity in this case) that determine the maximum subtree size with the same color. Hence, the flowchart clearly suggests opting for DFS given it's a tree-based problem.

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:

  1. Whether the current node a has the same color as its adjacent node b.
  2. 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.

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:

  1. Initialize a boolean variable ok to true. This variable signifies whether all nodes in the subtree rooted at node a have the same color.

  2. 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.

  3. 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.

  4. Additionally, during this traversal, we update size[a] by adding size[b] to accumulate the total size of all subtrees rooted at a.

  5. 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.

  6. 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.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

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

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

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

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

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

  1. 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.

  2. 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.

  3. 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

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.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

Which of the tree traversal order can be used to obtain elements in a binary search tree in sorted order?

Recommended Readings

Everything About Trees

Everything About Trees A tree is a type of graph data structure composed of nodes and edges Its main properties are It is acyclic doesn't contain any cycles There exists a path from the root to any node Has N 1 edges where N is the number of nodes in the tree and

DFS Fundamentals

https algomonster s3 us east 2 amazonaws com cover_photos dfs svg Depth First Search Prereqs Recursion Review problems recursion_intro Trees problems tree_intro With a solid understanding of recursion under our belts we are now ready to tackle one of the most useful techniques in coding interviews Depth First Search DFS

Dynamic Programming Introduction and Patterns

What is Dynamic Programming Prerequisite DFS problems dfs_intro Backtracking problems backtracking Memoization problems memoization_intro Pruning problems backtracking_pruning Dynamic programming is an algorithmic optimization technique that breaks down a complicated problem into smaller overlapping sub problems in a recursive manner and uses solutions to the sub problems to construct a solution

Want a Structured Path to Master System Design Too? Don’t Miss This!

Load More