2440. 创建价值相同的连通块

来源第 89 场双周赛 Q4难度困难分数2460

题目描述

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

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

你可以删除一些边，将这棵树分成几个连通块。一个连通块的价值定义为这个连通块中所有节点 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 * 10⁴
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 表示一棵合法的树。

解法

方法一：枚举连通块的个数

思考

删边后各连通块权和必须相等，故块数 \(k\) 必须整除总和 \(s\)，且单块目标 \(t=s/k\) 不小于最大点权。 \(n \le 2\times 10^4\)，从大到小枚举合法 \(k\)。

树上 DFS 累加子树和：等于 \(t\) 则视为切下并向上返回 \(0\)；超过 \(t\) 则失败。根返回 \(0\) 即划分成功，答案为 \(k-1\)。

假设连通块的个数为 \(k\)，那么要删除的边数为 \(k-1\)，每个连通块的价值为 \(\frac{s}{k}\)，其中 \(s\) 为 \(\textit{nums}\) 所有节点的值之和。

我们从大到小枚举 \(k\)，如果存在一个 \(k\)，使得 \(\frac{s}{k}\) 是整数，并且得到的每个连通块的价值都相等，那么直接返回 \(k-1\)。其中 \(k\) 的初始值为 \(\min(n, \frac{s}{mx})\)，记 \(mx\) 为 \(\textit{nums}\) 中的最大值。

关键点在于判断对于给定的 \(\frac{s}{k}\)，是否能划分出若干子树，使得每棵子树的价值都为 \(\frac{s}{k}\)。

这里我们通过 dfs 函数来判断，从上到下递归遍历求出各个子树的价值，如果子树价值和恰好为 \(\frac{s}{k}\)，说明此时划分成功，我们将价值置为 \(0\) 返回给上一层，表示此子树可以与父节点断开。如果子树价值之和大于 \(\frac{s}{k}\)，说明此时划分失败，我们返回 \(-1\)，表示无法划分。

时间复杂度 \(O(n \sqrt{s})\)。其中 \(n\) 和 \(s\) 分别为 \(\textit{nums}\) 的长度和 \(\textit{nums}\) 所有节点的值之和。

Python3JavaC++Go

class Solution:
    def componentValue(self, nums: List[int], edges: List[List[int]]) -> int:
        def dfs(i, fa):
            x = nums[i]
            for j in g[i]:
                if j != fa:
                    y = dfs(j, i)
                    if y == -1:
                        return -1
                    x += y
            if x > t:
                return -1
            return x if x < t else 0

        n = len(nums)
        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        s = sum(nums)
        mx = max(nums)
        for k in range(min(n, s // mx), 1, -1):
            if s % k == 0:
                t = s // k
                if dfs(0, -1) == 0:
                    return k - 1
        return 0

class Solution {
    private List<Integer>[] g;
    private int[] nums;
    private int t;

    public int componentValue(int[] nums, int[][] edges) {
        int n = nums.length;
        g = new List[n];
        this.nums = nums;
        Arrays.setAll(g, k -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        int s = sum(nums), mx = max(nums);
        for (int k = Math.min(n, s / mx); k > 1; --k) {
            if (s % k == 0) {
                t = s / k;
                if (dfs(0, -1) == 0) {
                    return k - 1;
                }
            }
        }
        return 0;
    }

    private int dfs(int i, int fa) {
        int x = nums[i];
        for (int j : g[i]) {
            if (j != fa) {
                int y = dfs(j, i);
                if (y == -1) {
                    return -1;
                }
                x += y;
            }
        }
        if (x > t) {
            return -1;
        }
        return x < t ? x : 0;
    }

    private int sum(int[] arr) {
        int x = 0;
        for (int v : arr) {
            x += v;
        }
        return x;
    }

    private int max(int[] arr) {
        int x = arr[0];
        for (int v : arr) {
            x = Math.max(x, v);
        }
        return x;
    }
}

class Solution {
public:
    int componentValue(vector<int>& nums, vector<vector<int>>& edges) {
        int n = nums.size();
        int s = accumulate(nums.begin(), nums.end(), 0);
        int mx = *max_element(nums.begin(), nums.end());
        int t = 0;
        unordered_map<int, vector<int>> g;
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        function<int(int, int)> dfs = [&](int i, int fa) -> int {
            int x = nums[i];
            for (int j : g[i]) {
                if (j != fa) {
                    int y = dfs(j, i);
                    if (y == -1) return -1;
                    x += y;
                }
            }
            if (x > t) return -1;
            return x < t ? x : 0;
        };
        for (int k = min(n, s / mx); k > 1; --k) {
            if (s % k == 0) {
                t = s / k;
                if (dfs(0, -1) == 0) {
                    return k - 1;
                }
            }
        }
        return 0;
    }
};

func componentValue(nums []int, edges [][]int) int {
    s, mx := 0, slices.Max(nums)
    for _, x := range nums {
        s += x
    }
    n := len(nums)
    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)
    }
    t := 0
    var dfs func(int, int) int
    dfs = func(i, fa int) int {
        x := nums[i]
        for _, j := range g[i] {
            if j != fa {
                y := dfs(j, i)
                if y == -1 {
                    return -1
                }
                x += y
            }
        }
        if x > t {
            return -1
        }
        if x < t {
            return x
        }
        return 0
    }
    for k := min(n, s/mx); k > 1; k-- {
        if s%k == 0 {
            t = s / k
            if dfs(0, -1) == 0 {
                return k - 1
            }
        }
    }
    return 0
}

方法二：枚举 + 显式栈

思考

删边之后每个连通块的价值和必须相同，因此块数 \(k\) 必须整除总和 \(s\)，目标值 \(t = s/k\) 也不能小于最大点权。 \(n \le 2 \times 10^4\)，从大到小枚举满足这两个条件的 \(k\) 是可行的；若对每个 \(k\) 都从根递归累加子树，链上的调用深度就是 \(n\)，会超出递归栈。

某个 \(k\) 能否划分，只取决于后序离开节点时子树还剩下多少价值。子节点都处理完以后，当前点权加上尚未切开的子树和：恰好等于 \(t\) 就记成 \(0\)，表示这条边可以删掉；小于 \(t\) 就把剩余和交给父节点；大于 \(t\) 则这个 \(k\) 失败。

显式栈按 \((\text{节点}, \text{父节点}, \text{状态})\) 模拟这次后序。进入节点时先压出栈标记、再压子节点，离开时按上面的规则写入剩余和。根的剩余和为 \(0\) 说明整棵树被切成若干价值和为 \(t\) 的块，此时删除的边数是 \(k - 1\)。