3004. 相同颜色的最大子树 🔒

难度中等

题目描述

给定一个二维整数数组 edges，表示一个有 n 个节点的树，节点编号从 0 到 n - 1，以节点 0 为根，其中 edges[i] = [u_i, v_i] 表示节点 v_i 和 u_i 之间存在一条边。

还给定一个 下标从 0 开始，大小为 n 的整数数组 colors，其中 colors[i] 表示节点 i 分配的颜色。

我们希望找到一个节点 v，使得 v 的子树中的每个节点具有相同的颜色。

返回 具有 尽可能多 节点且符合上述要求的子树大小。

示例 1：

输入：edges = [[0,1],[0,2],[0,3]], colors = [1,1,2,3]
输出：1
解释：每种颜色分别对应为：1 -> 红色，2 -> 绿色，3 -> 蓝色。我们可以看到以节点 0 为根的子树具有不同颜色的子节点。任何其他子树都是相同颜色的，并且大小为 1。因此，我们返回 1。

示例 2：

输入：edges = [[0,1],[0,2],[0,3]], colors = [1,1,1,1]
输出：4
解释：整个树具有相同的颜色，以节点 0 为根的子树具有节点数最多，为 4。因此，我们返回 4。

示例 3：

输入：edges = [[0,1],[0,2],[2,3],[2,4]], colors = [1,2,3,3,3]
输出：3
解释：每种颜色分别对应为：1 -> 红色，2 -> 绿色，3 -> 蓝色。我们可以看到以节点 0 为根的子树有不同颜色的子节点。其他任何子树都是相同颜色的，但以节点 2 为根的子树的大小为 3，这是最大的。因此，我们返回 3。

提示：

n == edges.length + 1
1 <= n <= 5 * 10⁴
edges[i] == [u_i, v_i]
0 <= u_i, v_i < n
colors.length == n
1 <= colors[i] <= 10⁵
输入被生成，使得由 edges 表示的图是一棵树。

解法

方法一：DFS

思考

树有 \(n \le 5 \times 10^4\) 个结点，枚举每个子树再检查颜色需要重复遍历。

一个子树同色，当且仅当根与每个儿子同色，且每个儿子的子树本身同色。子树规模可在后序中累加。

为此一次 DFS 同时返回布尔值并维护 \(\textit{size}\)。仅当整棵子树同色时，用其规模更新答案。

我们先根据题目给定的边的信息，构建一个邻接表 \(g\)，其中 \(g[a]\) 表示节点 \(a\) 的所有相邻节点。然后我们创建一个长度为 \(n\) 的数组 \(\textit{size}\)，其中 \(\textit{size}[a]\) 表示以节点 \(a\) 为根的子树的节点数。

接下来，我们设计一个函数 \(\textit{dfs}(a, fa)\)，它将返回以节点 \(a\) 为根的子树是否满足题目要求。函数 \(\textit{dfs}(a, fa)\) 的执行过程如下：

首先，我们用一个变量 \(ok\) 记录以节点 \(a\) 为根的子树是否满足题目要求，初始时 \(ok\) 为 \(\textit{true}\)。
接着，我们遍历节点 \(a\) 的所有相邻节点 \(b\)，如果 \(b\) 不是 \(a\) 的父节点 \(fa\)，那么我们递归调用 \(\textit{dfs}(b, a)\)，并将返回值保存到变量 \(t\) 中，并且更新 \(ok\) 为 \(ok\) 与 \(\textit{colors}[a] = \textit{colors}[b] \land t\) 的值，其中 \(\land\) 表示逻辑与运算。然后，我们更新 \(\textit{size}[a] = \textit{size}[a] + \textit{size}[b]\)。
然后，我们判断 \(ok\) 的值，如果 \(ok\) 为 \(\textit{true}\)，那么我们更新答案 \(\textit{ans} = \max(\textit{ans}, \textit{size}[a])\)。
最后，我们返回 \(ok\) 的值。

我们调用 \(\textit{dfs}(0, -1)\)，其中 \(0\) 表示根节点的编号， \(-1\) 表示根节点没有父节点。最终的答案即为 \(\textit{ans}\)。

时间复杂度 \(O(n)\)，空间复杂度 \(O(n)\)。其中 \(n\) 是节点的数量。

Python3JavaC++GoTypeScript

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

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;
    }
}

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;
    }
};

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
}

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;
}

方法二：显式栈

思考

树有 \(n \le 5 \times 10^4\) 个结点，枚举每个子树再检查颜色需要重复遍历。一条链上第一次递归总是走向唯一的孩子，深度就是 \(n\)，会超出默认递归上限。

一棵子树同色，当且仅当根与每个孩子同色，且每个孩子的子树本身同色。这个判断要等孩子都处理完才能做。

因此用显式栈做后序。进入结点时先压入离开标记，再压入孩子。离开时根据孩子是否同色更新当前子树，只在整棵子树同色时用其规模更新答案。