创建价值相同的连通块

标签: 深度优先搜索 数组 数学 枚举

难度: Hard

有一棵 n 个节点的无向树,节点编号为 0 到 n - 1 。

给你一个长度为 n 下标从 0 开始的整数数组 nums ,其中 nums[i] 表示第 i 个节点的值。同时给你一个长度为 n - 1 的二维整数数组 edges ,其中 edges[i] = [ai, bi] 表示节点 ai 与 bi 之间有一条边。

你可以 删除 一些边,将这棵树分成几个连通块。一个连通块的 价值 定义为这个连通块中 所有 节点 i 对应的 nums[i] 之和。

你需要删除一些边,删除后得到的各个连通块的价值都相等。请返回你可以删除的边数 最多 为多少。

示例 1:

输入:nums = [6,2,2,2,6], edges = [[0,1],[1,2],[1,3],[3,4]] 
输出:2 
解释:上图展示了我们可以删除边 [0,1] 和 [3,4] 。得到的连通块为 [0] ,[1,2,3] 和 [4] 。每个连通块的价值都为 6 。可以证明没有别的更好的删除方案存在了,所以答案为 2 。

示例 2:

输入:nums = [2], edges = []
输出:0
解释:没有任何边可以删除。

提示:

  • 1 <= n <= 2 * 104
  • nums.length == n
  • 1 <= nums[i] <= 50
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= edges[i][0], edges[i][1] <= n - 1
  • edges 表示一棵合法的树。

Submission

运行时间: 220 ms

内存: 32.7 MB

class Solution:
    def componentValue(self, nums: List[int], edges: List[List[int]]) -> int:
        if len(set(nums)) == 1:
            return len(nums) - 1
        g = [[] for _ in nums]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)
        
        
        def dfs(x: int, fa: int) -> int:
            s = nums[x]
            for y in g[x]:
                if y != fa:
                    res = dfs(y, x)
                    if res < 0:
                        return -1
                    s += res
            if s > target:
                return -1
            return s if s < target else 0

        total = sum(nums)
        for i in range(total // max(nums), 1, -1):
            if total % i == 0:
                target = total // i
                if dfs(0, -1) == 0: return i - 1
        return 0

Explain

该题解通过DFS和贪心算法的结合来解决问题。首先,题解构建了一个邻接列表表示的图来描述树。在解题的核心方法中,题解尝试将树分割成具有相同总和的多个连通块。这是通过尝试每个可能的连通块的目标和(即树节点值总和的因子)来完成的。对于每个目标和,使用DFS检查是否可以通过删除边来形成大小为目标和的连通块。如果从根节点开始的DFS返回0,表示可以形成这样的连通块。DFS函数在遍历时累积子树的值,如果子树的值等于目标和,则可以视为一个独立的连通块,因此该子树的贡献重置为0。如果子树的值超过了目标和,直接返回-1表示不可能形成所需的连通块。通过这种方式,题解实现了对可能删除的最大边数的检测。

时间复杂度: O(n * sqrt(total))

空间复杂度: O(n)

class Solution:
    def componentValue(self, nums: List[int], edges: List[List[int]]) -> int:
        # 特殊情况处理:如果所有节点的值相同,则最多可以删除所有边
        if len(set(nums)) == 1:
            return len(nums) - 1
        # 构建图
        g = [[] for _ in nums]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)
        
        # DFS函数,用于检查是否可以形成目标和为target的连通块
        def dfs(x: int, fa: int) -> int:
            s = nums[x]
            for y in g[x]:
                if y != fa:
                    res = dfs(y, x)
                    if res < 0:
                        return -1
                    s += res
            if s > target:
                return -1
            return s if s < target else 0

        total = sum(nums)
        # 尝试每个可能的目标和
        for i in range(total // max(nums), 1, -1):
            if total % i == 0:
                target = total // i
                if dfs(0, -1) == 0: return i - 1
        return 0

Explore

在这个题解中,DFS函数被用来递归地计算每个节点的子树总和。函数从某个节点开始,遍历其所有子节点,并累积这些子节点的子树之和。每次递归返回子树的总和,如果子树的总和达到了目标和(即我们想要的连通块的大小),则可以将这个子树视作一个独立的连通块,并通过返回0而将这部分子树的贡献'重置',这样上层节点就不会再包括这部分的总和,有助于进一步的连通块的形成。如果某个子树的总和超过了目标和,则返回-1表示当前的分割方式不可行。这样,通过DFS的递归调用,能够检查是否可以通过删除某些边来形成多个总和等于目标和的连通块。

当DFS过程中子树的总和等于目标和时,将返回值设为0的处理方式主要是为了表示这部分子树已经可以形成一个独立的连通块,因此它的贡献应从其父节点的累积总和中移除,这样父节点及其他子树就可以继续寻找或构建其他的连通块。这种处理方式有效地帮助我们标记已经成功形成的连通块,并防止它们被重复计算或错误地划分到其他连通块中。

如果在DFS过程中某个子树的值超过了目标和而返回-1,这通常表示当前的分割路径不可行,但并不意味着整个分割尝试失败。返回-1仅代表从当前节点出发的这一特定分割方式无法成功形成目标和大小的连通块。算法会继续探索其他可能的节点和分割方式,直到找到可行的解或遍历完所有可能性。

在DFS函数中,参数`fa`代表当前节点的父节点。这个参数的主要作用是防止DFS过程中的递归调用回溯到上一级节点,从而避免形成无限循环和错误的累积总和。在树或无向图的DFS遍历中,标记或记录父节点是常见的技术,用于确保每个节点只被访问一次,从而正确地计算出每个节点的子树总和。