Skip to content

3004. Maximum Subtree of the Same Color πŸ”’

DifficultyMedium

Description

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] = [ui, vi] means there is an edge between the nodes vi and ui.

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 subtree with the maximum number of nodes possible.

 

Example 1:

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.

Example 2:

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.

Example 3:

Input: edges = [[0,1],[0,2],[2,3],[2,4]], colors = [1,2,3,3,3]
Output: 3
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, but the subtree rooted at node 2 has a size of 3 which is the maximum. Hence, we return 3.

 

Constraints:

  • n == edges.length + 1
  • 1 <= n <= 5 * 104
  • edges[i] == [ui, vi]
  • 0 <= ui, vi < n
  • colors.length == n
  • 1 <= colors[i] <= 105
  • The input is generated such that the graph represented by edges is a tree.

Solutions

Solution 1: DFS

Thinking

The tree has \(n \le 5 \times 10^4\) nodes, so rechecking every subtree for a single color would repeat work.

A subtree is monochromatic iff the root matches every child and every child’s subtree is itself monochromatic. Sizes add up in post-order.

One DFS therefore returns a boolean and maintains \(\textit{size}\). Only a fully monochromatic subtree updates the answer with its size.

First, according to the edge information given in the problem, we construct an adjacency list \(g\), where \(g[a]\) represents all adjacent nodes of node \(a\). Then we create an array \(\textit{size}\) of length \(n\), where \(\textit{size}[a]\) represents the number of nodes in the subtree with node \(a\) as the root.

Next, we design a function \(\textit{dfs}(a, fa)\), which will return whether the subtree with node \(a\) as the root meets the requirements of the problem. The execution process of the function \(\textit{dfs}(a, fa)\) is as follows:

  • First, we use a variable \(ok\) to record whether the subtree with node \(a\) as the root meets the requirements of the problem, and initially \(ok\) is \(\textit{true}\).
  • Then, we traverse all adjacent nodes \(b\) of node \(a\). If \(b\) is not the parent node \(fa\) of \(a\), then we recursively call \(\textit{dfs}(b, a)\), save the return value into the variable \(t\), and update \(ok\) to the value of \(ok\) and \(\textit{colors}[a] = \textit{colors}[b] \land t\), where \(\land\) represents logical AND operation. Then, we update \(\textit{size}[a] = \textit{size}[a] + \textit{size}[b]\).
  • After that, we judge the value of \(ok\). If \(ok\) is \(\textit{true}\), then we update the answer \(\textit{ans} = \max(\textit{ans}, \textit{size}[a])\).
  • Finally, we return the value of \(ok\).

We call \(\textit{dfs}(0, -1)\), where \(0\) represents the root node number, and \(-1\) represents that the root node has no parent node. The final answer is \(\textit{ans}\).

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the number of nodes.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
    private List<Integer>[] g;
    private int[] colors;
    private int[] size;
    private int ans;

    public int maximumSubtreeSize(int[][] edges, int[] colors) {
        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;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Solution {
public:
    int maximumSubtreeSize(vector<vector<int>>& edges, vector<int>& colors) {
        int n = edges.size() + 1;
        vector<int> g[n];
        vector<int> size(n, 1);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        int ans = 0;
        function<bool(int, int)> dfs = [&](int a, int fa) {
            bool ok = true;
            for (int b : g[a]) {
                if (b != fa) {
                    bool t = dfs(b, a);
                    ok = ok && colors[a] == colors[b] && t;
                    size[a] += size[b];
                }
            }
            if (ok) {
                ans = max(ans, size[a]);
            }
            return ok;
        };
        dfs(0, -1);
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
func maximumSubtreeSize(edges [][]int, colors []int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    size := make([]int, n)
    var dfs func(int, int) bool
    dfs = func(a, fa int) bool {
        size[a] = 1
        ok := true
        for _, b := range g[a] {
            if b != fa {
                t := dfs(b, a)
                ok = ok && t && colors[a] == colors[b]
                size[a] += size[b]
            }
        }
        if ok {
            ans = max(ans, size[a])
        }
        return ok
    }
    dfs(0, -1)
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
function maximumSubtreeSize(edges: number[][], colors: number[]): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const size: number[] = Array(n).fill(1);
    let ans = 0;
    const dfs = (a: number, fa: number): boolean => {
        let ok = true;
        for (const b of g[a]) {
            if (b !== fa) {
                const t = dfs(b, a);
                ok = ok && t && colors[a] === colors[b];
                size[a] += size[b];
            }
        }
        if (ok) {
            ans = Math.max(ans, size[a]);
        }
        return ok;
    };
    dfs(0, -1);
    return ans;
}

Solution 2: Explicit Stack

Thinking

The tree has \(n \le 5 \times 10^4\) nodes, so rechecking every subtree for a single color would repeat work. On a chain the first recursive call always follows the only child, so the depth is \(n\) and exceeds the default recursion limit.

A subtree is monochromatic only when the root matches every child and every child's subtree is itself monochromatic. That test waits until the children are finished.

An explicit stack walks the tree in postorder. On entry, push the exit marker and then the children. On exit, combine the children's colors and sizes, and update the answer only when the whole subtree is one color.

First, according to the edge information given in the problem, we construct an adjacency list \(g\), where \(g[a]\) represents all adjacent nodes of node \(a\). Then we create an array \(\textit{size}\) of length \(n\), where \(\textit{size}[a]\) represents the number of nodes in the subtree with node \(a\) as the root.

An explicit stack walks the tree from root \(0\) in postorder. Entering a node pushes that node's exit marker and then its children, so the children finish first. Every node's \(\textit{size}\) starts at \(1\). On exit, \(ok\) records whether the subtree is monochromatic and starts true. For each child \(b\), \(ok\) becomes the conjunction of its current value, \(\textit{colors}[a] = \textit{colors}[b]\), and the child's own monochromatic flag, and \(\textit{size}[b]\) is added to \(\textit{size}[a]\). If \(ok\) is still true, \(\textit{size}[a]\) updates the answer.

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the number of nodes.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution:
    def maximumSubtreeSize(self, edges: List[List[int]], colors: List[int]) -> int:
        n = len(edges) + 1
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        size = [1] * n
        ok = [False] * n
        ans = 0
        stk = [(0, -1, 0)]
        while stk:
            a, fa, state = stk.pop()
            if state == 0:
                stk.append((a, fa, 1))
                for b in g[a]:
                    if b != fa:
                        stk.append((b, a, 0))
            else:
                good = True
                for b in g[a]:
                    if b != fa:
                        good = good and colors[a] == colors[b] and ok[b]
                        size[a] += size[b]
                if good:
                    ans = max(ans, size[a])
                ok[a] = good
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
    public int maximumSubtreeSize(int[][] edges, int[] colors) {
        int n = edges.length + 1;
        List<Integer>[] g = new List[n];
        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);
        }
        int[] size = new int[n];
        Arrays.fill(size, 1);
        boolean[] ok = new boolean[n];
        int ans = 0;
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int a = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push(new int[] {a, fa, 1});
                for (int b : g[a]) {
                    if (b != fa) {
                        stk.push(new int[] {b, a, 0});
                    }
                }
            } else {
                boolean good = true;
                for (int b : g[a]) {
                    if (b != fa) {
                        good = good && colors[a] == colors[b] && ok[b];
                        size[a] += size[b];
                    }
                }
                if (good) {
                    ans = Math.max(ans, size[a]);
                }
                ok[a] = good;
            }
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution {
public:
    int maximumSubtreeSize(vector<vector<int>>& edges, vector<int>& colors) {
        int n = edges.size() + 1;
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        vector<int> size(n, 1);
        vector<char> ok(n);
        int ans = 0;
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto cur = stk.back();
            stk.pop_back();
            int a = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push_back({a, fa, 1});
                for (int b : g[a]) {
                    if (b != fa) {
                        stk.push_back({b, a, 0});
                    }
                }
            } else {
                bool good = true;
                for (int b : g[a]) {
                    if (b != fa) {
                        good = good && colors[a] == colors[b] && ok[b];
                        size[a] += size[b];
                    }
                }
                if (good) {
                    ans = max(ans, size[a]);
                }
                ok[a] = good;
            }
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
func maximumSubtreeSize(edges [][]int, colors []int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    size := make([]int, n)
    for i := range size {
        size[i] = 1
    }
    ok := make([]bool, n)
    stk := [][3]int{{0, -1, 0}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        a, fa, state := cur[0], cur[1], cur[2]
        if state == 0 {
            stk = append(stk, [3]int{a, fa, 1})
            for _, b := range g[a] {
                if b != fa {
                    stk = append(stk, [3]int{b, a, 0})
                }
            }
        } else {
            good := true
            for _, b := range g[a] {
                if b != fa {
                    good = good && colors[a] == colors[b] && ok[b]
                    size[a] += size[b]
                }
            }
            if good {
                ans = max(ans, size[a])
            }
            ok[a] = good
        }
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
function maximumSubtreeSize(edges: number[][], colors: number[]): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const size: number[] = Array(n).fill(1);
    const ok: boolean[] = Array(n).fill(false);
    let ans = 0;
    const stk: number[][] = [[0, -1, 0]];
    while (stk.length) {
        const [a, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([a, fa, 1]);
            for (const b of g[a]) {
                if (b !== fa) {
                    stk.push([b, a, 0]);
                }
            }
        } else {
            let good = true;
            for (const b of g[a]) {
                if (b !== fa) {
                    good = good && colors[a] === colors[b] && ok[b];
                    size[a] += size[b];
                }
            }
            if (good) {
                ans = Math.max(ans, size[a]);
            }
            ok[a] = good;
        }
    }
    return ans;
}

Comments