Maximum subtree of the same color

Home | Coding Interviews | Graphs | Maximum subtree of the same color

Graphs Medium Video Python Code Solution Java

You are 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, v_i] means there is an edge between the nodes v_i and u_i.

You are also given a 0-indexed integer array colors of size n, where colors[i] is the color assigned to node i.

We want to find a node v such that every node in the subtree of v has the same color.

Return the size of such a subtree with the maximum number of nodes possible.

class Solution {
    private List<Integer>[] g;
    private int[] colors;
    private int[] size;
    private int ans;

    public int maximumSubtreeSize(int[][] edges, int[] colors) {//build adjacency list
        int n = edges.length + 1;
        g = new List[n];
        size = new int[n];
        this.colors = colors;
        Arrays.fill(size, 1);
        Arrays.setAll(g, i -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        dfs(0, -1);
        return ans;
    }

    private boolean dfs(int a, int fa) {
        boolean ok = true;
        for (int b : g[a]) {
            if (b != fa) {
                boolean t = dfs(b, a);
                ok = ok && colors[a] == colors[b] && t;
                size[a] += size[b];
            }
        }
        if (ok) {
            ans = Math.max(ans, size[a]);
        }
        return ok;
    }
}

class Solution:
    def maximumSubtreeSize(self, edges: List[List[int]], colors: List[int]) -> int:
        def dfs(a: int, fa: int) -> bool:
            ok = True
            for b in g[a]:
                if b != fa:
                    t = dfs(b, a)
                    ok = ok and colors[a] == colors[b] and t
                    size[a] += size[b]
            if ok:
                nonlocal ans
                ans = max(ans, size[a])
            return ok

        n = len(edges) + 1
        g = [[] for _ in range(n)]
        size = [1] * n
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = 0
        dfs(0, -1)
        return ans

To solve this problem using Depth-First Search (DFS), traverse the tree starting from the root node. During the traversal, keep track of the colors of nodes in each subtree. For each node, recursively determine if all nodes in its subtree have the same color. If they do, compute the size of this subtree and update the maximum size found.

Specifically, implement a DFS function that returns both the size of the subtree and a boolean indicating whether all nodes in the subtree have the same color. If the subtree rooted at a node satisfies the condition, compare its size with the current maximum and update if larger. This approach ensures an efficient traversal and comparison, yielding the desired result.

Posted by grwgreg 2 months ago

Maximum subtree of the same color

Related Problems