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:

C++

#include <bits/stdc++.h>
using namespace std;

// This variable stores the maximum size of a subtree that
// has all nodes of the same color.
int m_maxSize = 1;

// This struct represents the result of a depth-first search
// (DFS) on a subtree.
struct Res {
    // This variable indicates whether all nodes in the
    // subtree have the same color.
    bool all_same_color;

    // This variable stores the number of nodes in the
    // subtree.
    int num_nodes;

    // This variable stores the color of the nodes in the
    // subtree.
    int color;
};

// This function performs a DFS on a subtree and returns the
// result.
Res dfs(int node, unordered_map<int, vector<int> >& graph,
        vector<int>& colors)
{
    // If the node has no children, then it is a leaf node
    // and the result is simply that all nodes have the same
    // color (true), the number of nodes is 1, and the color
    // is the color of the node.
    if (graph[node].empty()) {
        return { true, 1, colors[node] };
    }

    // This variable indicates whether all nodes in the
    // subtree have the same color.
    bool color_match = true;

    // This variable stores the sum of the number of nodes
    // in the subtree.
    int sum_of_nodes = 0;

    // This loop iterates over the children of the node.
    for (int child : graph[node]) {
        // This variable stores the result of the DFS on the
        // child subtree.
        Res res = dfs(child, graph, colors);

        // This variable indicates whether all nodes in the
        // child subtree have the same color.
        bool all_same_color = res.all_same_color;

        // This variable stores the number of nodes in the
        // child subtree.
        int num_nodes = res.num_nodes;

        // This variable stores the color of the nodes in
        // the child subtree.
        int color = res.color;

        // This variable indicates whether the color of the
        // node matches the color of the child subtree.
        color_match = color_match && all_same_color;

        // If the color of the node does not match the color
        // of the child subtree, then the color_match
        // variable is set to false.
        if (color != colors[node]) {
            color_match = false;
        }

        // This variable stores the sum of the number of
        // nodes in the subtree.
        sum_of_nodes += num_nodes;
    }

    // If all nodes in the subtree have the same color, then
    // the maximum size of the subtree is updated.
    if (color_match) {
        m_maxSize = max(m_maxSize, sum_of_nodes + 1);
    }

    // This variable stores the result of the DFS on the
    // subtree.
    return { color_match, sum_of_nodes + 1, colors[node] };
}

// This function returns the maximum size of a subtree that
// has all nodes of the same color.
int maximumSubtreeSize(vector<vector<int> >& edges,
                       vector<int>& colors)
{
    // This variable stores the graph as an adjacency list.
    unordered_map<int, vector<int> > graph;

    // This loop iterates over the edges and adds them to
    // the graph.
    for (auto edge : edges) {
        graph[edge[0]].push_back(edge[1]);
    }

    // This variable stores the result of the DFS on the
    // entire graph.
    dfs(0, graph, colors);

    // This variable returns the maximum size of a subtree
    // that has all nodes of the same color.
    return m_maxSize;
}

// This is the driver code.
int main()
{
    // This variable stores the edges of the graph.
    vector<vector<int> > edges = {
        { 0, 1 }, { 0, 2 }, { 1, 3 }, { 1, 4 }, { 2, 5 }
    };

    // This variable stores the colors of the nodes in the
    // graph.
    vector<int> colors = { 1, 1, 2, 1, 2, 2 };

    // This variable stores the maximum size of a subtree
    // that has all nodes of the same color.
    int maximum_subtree_size
        = maximumSubtreeSize(edges, colors);

    // This variable prints the maximum size of a subtree
    // that has all nodes of the same color.
    cout << maximum_subtree_size << endl;

    return 0;
}

Java

import java.util.*;

public class Main {
    // This variable stores the maximum size of a subtree that
    // has all nodes of the same color.
    static int maxSize = 1;

    // This class represents the result of a depth-first search
    // (DFS) on a subtree.
    static class Res {
        // This variable indicates whether all nodes in the
        // subtree have the same color.
        boolean allSameColor;

        // This variable stores the number of nodes in the
        // subtree.
        int numNodes;

        // This variable stores the color of the nodes in the
        // subtree.
        int color;

        Res(boolean allSameColor, int numNodes, int color) {
            this.allSameColor = allSameColor;
            this.numNodes = numNodes;
            this.color = color;
        }
    }

    // This function performs a DFS on a subtree and returns the
    // result.
    static Res dfs(int node, Map<Integer, List<Integer>> graph,
                   List<Integer> colors) {
        // If the node has no children, then it is a leaf node
        // and the result is simply that all nodes have the same
        // color (true), the number of nodes is 1, and the color
        // is the color of the node.
        if (!graph.containsKey(node) || graph.get(node).isEmpty()) {
            return new Res(true, 1, colors.get(node));
        }

        // This variable indicates whether all nodes in the
        // subtree have the same color.
        boolean colorMatch = true;

        // This variable stores the sum of the number of nodes
        // in the subtree.
        int sumOfNodes = 0;

        // This loop iterates over the children of the node.
        for (int child : graph.get(node)) {
            // This variable stores the result of the DFS on the
            // child subtree.
            Res res = dfs(child, graph, colors);

            // This variable indicates whether all nodes in the
            // child subtree have the same color.
            boolean allSameColor = res.allSameColor;

            // This variable stores the number of nodes in the
            // child subtree.
            int numNodes = res.numNodes;

            // This variable stores the color of the nodes in
            // the child subtree.
            int color = res.color;

            // This variable indicates whether the color of the
            // node matches the color of the child subtree.
            colorMatch = colorMatch && allSameColor;

            // If the color of the node does not match the color
            // of the child subtree, then the colorMatch
            // variable is set to false.
            if (color != colors.get(node)) {
                colorMatch = false;
            }

            // This variable stores the sum of the number of
            // nodes in the subtree.
            sumOfNodes += numNodes;
        }

        // If all nodes in the subtree have the same color, then
        // the maximum size of the subtree is updated.
        if (colorMatch) {
            maxSize = Math.max(maxSize, sumOfNodes + 1);
        }

        // This variable stores the result of the DFS on the
        // subtree.
        return new Res(colorMatch, sumOfNodes + 1, colors.get(node));
    }

    // This function returns the maximum size of a subtree that
    // has all nodes of the same color.
    static int maximumSubtreeSize(List<List<Integer>> edges,
                                   List<Integer> colors) {
        // This variable stores the graph as an adjacency list.
        Map<Integer, List<Integer>> graph = new HashMap<>();

        // This loop iterates over the edges and adds them to
        // the graph.
        for (List<Integer> edge : edges) {
            int u = edge.get(0), v = edge.get(1);
            graph.computeIfAbsent(u, k -> new ArrayList<>()).add(v);
        }

        // This variable stores the result of the DFS on the
        // entire graph.
        dfs(0, graph, colors);

        // This variable returns the maximum size of a subtree
        // that has all nodes of the same color.
        return maxSize;
    }

    // This is the driver code.
    public static void main(String[] args) {
        // This variable stores the edges of the graph.
        List<List<Integer>> edges = Arrays.asList(
                Arrays.asList(0, 1),
                Arrays.asList(0, 2),
                Arrays.asList(1, 3),
                Arrays.asList(1, 4),
                Arrays.asList(2, 5)
        );

        // This variable stores the colors of the nodes in the
        // graph.
        List<Integer> colors = Arrays.asList(1, 1, 2, 1, 2, 2);

        // This variable stores the maximum size of a subtree
        // that has all nodes of the same color.
        int maximumSubtreeSize = maximumSubtreeSize(edges, colors);

        // This variable prints the maximum size of a subtree
        // that has all nodes of the same color.
        System.out.println(maximumSubtreeSize);
    }
}

// This code is contributed by shivamgupta310570

Python3

# This variable stores the maximum size of a subtree that
# has all nodes of the same color.
m_maxSize = 1

# This struct represents the result of a depth-first search
# (DFS) on a subtree.
class Res:
    def __init__(self, all_same_color, num_nodes, color):
        self.all_same_color = all_same_color
        self.num_nodes = num_nodes
        self.color = color

# This function performs a DFS on a subtree and returns the
# result.
def dfs(node, graph, colors):
    # If the node has no children, then it is a leaf node
    # and the result is simply that all nodes have the same
    # color (true), the number of nodes is 1, and the color
    # is the color of the node.
    if node not in graph:
        return Res(True, 1, colors[node])

    # This variable indicates whether all nodes in the
    # subtree have the same color.
    color_match = True

    # This variable stores the sum of the number of nodes
    # in the subtree.
    sum_of_nodes = 0

    # This loop iterates over the children of the node.
    for child in graph[node]:
        # This variable stores the result of the DFS on the
        # child subtree.
        res = dfs(child, graph, colors)

        # This variable indicates whether all nodes in the
        # child subtree have the same color.
        all_same_color = res.all_same_color

        # This variable stores the number of nodes in the
        # child subtree.
        num_nodes = res.num_nodes

        # This variable stores the color of the nodes in
        # the child subtree.
        color = res.color

        # This variable indicates whether the color of the
        # node matches the color of the child subtree.
        color_match = color_match and all_same_color

        # If the color of the node does not match the color
        # of the child subtree, then the color_match
        # variable is set to false.
        if color != colors[node]:
            color_match = False

        # This variable stores the sum of the number of
        # nodes in the subtree.
        sum_of_nodes += num_nodes

    # If all nodes in the subtree have the same color, then
    # the maximum size of the subtree is updated.
    if color_match:
        global m_maxSize
        m_maxSize = max(m_maxSize, sum_of_nodes + 1)

    # This variable stores the result of the DFS on the
    # subtree.
    return Res(color_match, sum_of_nodes + 1, colors[node])

# This function returns the maximum size of a subtree that
# has all nodes of the same color.
def maximumSubtreeSize(edges, colors):
    # This variable stores the graph as an adjacency list.
    graph = {}

    # This loop iterates over the edges and adds them to
    # the graph.
    for edge in edges:
        if edge[0] not in graph:
            graph[edge[0]] = []
        graph[edge[0]].append(edge[1])

    # This variable stores the result of the DFS on the
    # entire graph.
    dfs(0, graph, colors)

    # This variable returns the maximum size of a subtree
    # that has all nodes of the same color.
    return m_maxSize

# This is the driver code.
if __name__ == "__main__":
    # This variable stores the edges of the graph.
    edges = [
        [0, 1], [0, 2], [1, 3], [1, 4], [2, 5]
    ]

    # This variable stores the colors of the nodes in the
    # graph.
    colors = [1, 1, 2, 1, 2, 2]

    # This variable stores the maximum size of a subtree
    # that has all nodes of the same color.
    maximum_subtree_size = maximumSubtreeSize(edges, colors)

    # This variable prints the maximum size of a subtree
    # that has all nodes of the same color.
    print(maximum_subtree_size)

    # This code is contributed by Ayush Mishra

JavaScript

// This variable stores the maximum size of a subtree that
// has all nodes of the same color.
let m_maxSize = 1;

// This class represents the result of a depth-first search
// (DFS) on a subtree.
class Res {
    constructor(all_same_color, num_nodes, color) {
        this.all_same_color = all_same_color;
        this.num_nodes = num_nodes;
        this.color = color;
    }
}

// This function performs a DFS on a subtree and returns the result.
function dfs(node, graph, colors) {
    // If the node has no children, then it is a leaf node
    // and the result is simply that all nodes have the same
    // color (true), the number of nodes is 1, and the color
    // is the color of the node.
    if (!(node in graph)) {
        return new Res(true, 1, colors[node]);
    }

    // This variable indicates whether all nodes in the
    // subtree have the same color.
    let color_match = true;

    // This variable stores the sum of the number of nodes
    // in the subtree.
    let sum_of_nodes = 0;

    // This loop iterates over the children of the node.
    for (const child of graph[node]) {
        // This variable stores the result of the DFS on the
        // child subtree.
        const res = dfs(child, graph, colors);

        // This variable indicates whether all nodes in the
        // child subtree have the same color.
        const all_same_color = res.all_same_color;

        // This variable stores the number of nodes in the
        // child subtree.
        const num_nodes = res.num_nodes;

        // This variable stores the color of the nodes in
        // the child subtree.
        const color = res.color;

        // This variable indicates whether the color of the
        // node matches the color of the child subtree.
        color_match = color_match && all_same_color;

        // If the color of the node does not match the color
        // of the child subtree, then the color_match
        // variable is set to false.
        if (color !== colors[node]) {
            color_match = false;
        }

        // This variable stores the sum of the number of
        // nodes in the subtree.
        sum_of_nodes += num_nodes;
    }

    // If all nodes in the subtree have the same color, then
    // the maximum size of the subtree is updated.
    if (color_match) {
        m_maxSize = Math.max(m_maxSize, sum_of_nodes + 1);
    }

    // This variable stores the result of the DFS on the
    // subtree.
    return new Res(color_match, sum_of_nodes + 1, colors[node]);
}

// This function returns the maximum size of a subtree that
// has all nodes of the same color.
function maximumSubtreeSize(edges, colors) {
    // This variable stores the graph as an adjacency list.
    const graph = {};

    // This loop iterates over the edges and adds them to
    // the graph.
    for (const [u, v] of edges) {
        if (!(u in graph)) {
            graph[u] = [];
        }
        graph[u].push(v);
    }

    // This variable stores the result of the DFS on the
    // entire graph.
    dfs(0, graph, colors);

    // This variable returns the maximum size of a subtree
    // that has all nodes of the same color.
    return m_maxSize;
}

// This is the driver code.
const edges = [
    [0, 1], [0, 2], [1, 3], [1, 4], [2, 5]
];

const colors = [1, 1, 2, 1, 2, 2];

// This variable stores the maximum size of a subtree
// that has all nodes of the same color.
const maximum_subtree_size = maximumSubtreeSize(edges, colors);

// This variable prints the maximum size of a subtree
// that has all nodes of the same color.
console.log(maximum_subtree_size);

Output

2

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