可以被 K 整除连通块的最大数目

难度: Hard

给你一棵 n 个节点的无向树，节点编号为 0 到 n - 1 。给你整数 n 和一个长度为 n - 1 的二维整数数组 edges ，其中 edges[i] = [a_i, b_i] 表示树中节点 a_i 和 b_i 有一条边。

同时给你一个下标从 0 开始长度为 n 的整数数组 values ，其中 values[i] 是第 i 个节点的值。再给你一个整数 k 。

你可以从树中删除一些边，也可以一条边也不删，得到若干连通块。一个 连通块的值 定义为连通块中所有节点值之和。如果所有连通块的值都可以被 k 整除，那么我们说这是一个 合法分割 。

请你返回所有合法分割中，连通块数目的最大值 。

示例 1：

输入：n = 5, edges = [[0,2],[1,2],[1,3],[2,4]], values = [1,8,1,4,4], k = 6
输出：2
解释：我们删除节点 1 和 2 之间的边。这是一个合法分割，因为：
- 节点 1 和 3 所在连通块的值为 values[1] + values[3] = 12 。
- 节点 0 ，2 和 4 所在连通块的值为 values[0] + values[2] + values[4] = 6 。
最多可以得到 2 个连通块的合法分割。

示例 2：

输入：n = 7, edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], values = [3,0,6,1,5,2,1], k = 3
输出：3
解释：我们删除节点 0 和 2 ，以及节点 0 和 1 之间的边。这是一个合法分割，因为：
- 节点 0 的连通块的值为 values[0] = 3 。
- 节点 2 ，5 和 6 所在连通块的值为 values[2] + values[5] + values[6] = 9 。
- 节点 1 ，3 和 4 的连通块的值为 values[1] + values[3] + values[4] = 6 。
最多可以得到 3 个连通块的合法分割。

提示：

1 <= n <= 3 * 10⁴
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
values.length == n
0 <= values[i] <= 10⁹
1 <= k <= 10⁹
values 之和可以被 k 整除。
输入保证 edges 是一棵无向树。

Submission

运行时间: 173 ms

内存: 34.9 MB

class Solution:
    def maxKDivisibleComponents(self, n: int, edges: List[List[int]], values: List[int], k: int) -> int:

        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)
        ans = 0

        def dfs(x, fa):
            v = values[x]
            for y in g[x]:
                if y != fa:
                    v += dfs(y, x)
            nonlocal ans
            if v % k == 0:
                ans += 1
            return v
        
        dfs(0, -1)
        return ans

Explain

题解采用了深度优先搜索（DFS）来遍历树，并计算各个子树的节点值之和。在遍历过程中，每当一个子树的节点值之和可以被k整除时，就将结果计数器（ans）加1。这样，计数器最终的值就是可以分割成的最大连通块数。函数dfs(x, fa)用于从节点x开始，避开父节点fa，递归地计算子树的值之和。如果子树的值之和可以被k整除，则意味着可以在此处断开连接，从而形成一个新的连通块。

时间复杂度: O(n)

空间复杂度: O(n)

# 定义解决方案类

class Solution:
    def maxKDivisibleComponents(self, n: int, edges: List[List[int]], values: List[int], k: int) -> int:
        # 创建图的邻接表
        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)
        # 初始化答案变量
        ans = 0

        # 定义深度优先搜索函数
        def dfs(x, fa):
            # 初始化当前节点的值
            v = values[x]
            # 遍历当前节点的所有邻接节点
            for y in g[x]:
                if y != fa:  # 避免访问父节点
                    v += dfs(y, x)  # 递归计算子树的值之和
            # 使用nonlocal声明外部变量
            nonlocal ans
            # 如果当前子树的值之和可以被k整除
            if v % k == 0:
                ans += 1  # 增加答案计数
            return v  # 返回当前子树的值之和
        
        # 从节点0开始DFS遍历
        dfs(0, -1)
        # 返回结果，需要减去整棵树本身也是一个连通块的情况
        return ans

Explore

在计算连通块的值之和时，`values`数组中的负数或零值按照其算术值正常处理。这意味着，负数会减少总和，而零值不会影响总和。在求和的过程中，只要最终的子树值之和能被`k`整除，无论是由于包含负数、零值还是正数，都会被视为有效的连通块。因此，这些特殊值不会影响连通块计数的逻辑，仅影响计算的数值结果。

在题解的实现中，如果整棵树的值之和也可以被`k`整除，那么整棵树本身也被错误地计为一个可以分割的连通块。为了纠正这一点，我们需要从最终的结果中减去1。这是因为题目要求的是可以分割的连通块的最大数目，而整棵树本身不能被视为分割出的连通块。因此，在返回答案之前应减去1，确保不重复计算整棵树。

在DFS递归过程中立即增加答案计数允许我们在每次递归调用返回时就确定是否形成了一个新的连通块。这种方法利用了递归的自然结构，即从底部到顶部逐步构建解决方案的过程。如果等到整个递归完成后再统一处理，我们将需要额外的数据结构来存储每个子树的计算结果，这样会增加算法的复杂性和可能的错误源。因此，贪心地在每次符合条件的递归返回时更新结果是更直接和高效的处理方式。

对于非常大的树，特别是当树的结构呈现为一条链时，递归的深度优先搜索可能会导致深度递归栈溢出。为了优化这种情况，可以考虑使用迭代的深度优先搜索（使用显式栈）或广度优先搜索（BFS）。这些迭代方法避免了使用递归调用栈，从而减少了栈溢出的风险。此外，对于特别大的数据结构，可以考虑使用尾递归优化，虽然这在Python中不是自动的，但可以通过手动优化递归逻辑来实现。