2378. 选择边来最大化树的得分 🔒

难度中等

题目描述

给定一个加权树，由 n 个节点组成，从 0 到 n - 1。

该树以节点 0 为根，用大小为 n 的二维数组 edges 表示，其中 edges[i] = [par_i, weight_i] 表示节点 par_i 是节点 i 的父节点，它们之间的边的权重等于 weight_i。因为根结点没有父结点，所以有 edges[0] = [-1, -1]。

从树中选择一些边，使所选的两条边都不相邻，所选边的权值之和最大。

返回所选边的最大和。

注意:

你可以 不选择 树中的任何边，在这种情况下权值和将为 0。
如果树中的两条边 Edge₁ 和 Edge₂ 有一个公共节点，它们就是相邻的。
- 换句话说，如果 Edge₁连接节点 a 和 b, Edge₂ 连接节点 b 和 c，它们是相邻的。

示例 1:

输入: edges = [[-1,-1],[0,5],[0,10],[2,6],[2,4]]
输出: 11
解释: 上面的图表显示了我们必须选择红色的边。
总分是 5 + 6 = 11.
可以看出，没有更好的分数可以获得。

示例 2:

输入: edges = [[-1,-1],[0,5],[0,-6],[0,7]]
输出: 7
解释: 我们选择权值为 7 的边。
注意，我们不能选择一条以上的边，因为所有的边都是彼此相邻的。

提示:

n == edges.length
1 <= n <= 10⁵
edges[i].length == 2
par₀ == weight₀ == -1
i >= 1 时 0 <= par_i <= n - 1 。
par_i != i
i >= 1 时 -10⁶ <= weight_i <= 10⁶ 。
edges 表示有效的树。

解法

方法一：树形 DP

思考

在树中选一组不相邻的边使权和最大。 \(n\) 与树同阶时，选或不选一条边只影响其两端。

\(\textit{dfs}(i)\) 返回二元组：连向父边被选时的最大权和，以及不被选时的最大权和。前者只能累加儿子「不选父边」的值；后者可在至多一个儿子上改选「选父边」并加上该边权。

我们设计一个函数 \(\textit{dfs}(i)\)，表示以节点 \(i\) 为根的子树中，选择一些边，使得所选的两条边都不相邻，所选边的权值之和最大。该函数返回了两个值 \((a, b)\)，第一个值 \(a\) 表示当前节点 \(i\) 与其父节点之间的边被选中时，所选边的权值之和；第二个值 \(b\) 表示当前节点 \(i\) 与其父节点之间的边不被选中时，所选边的权值之和。

我们可以发现，对于当前节点 \(i\)：

如果 \(i\) 与父节点的边被选择，则它与子节点的所有边都不能被选择，那么当前节点的 \(a\) 值就是其所有子节点的 \(b\) 值之和；
如果 \(i\) 与父节点的边没被选择，那么可以选择它与子节点的最多一条边，那么当前节点的 \(b\) 值就是其选中的子节点的 \(a\) 值与未选中的子节点的 \(b\) 值之和，再加上 \(i\) 与选中的子节点之间的边的权值。

我们调用 \(\textit{dfs}(0)\) 函数，返回的第二个值即为答案，即当前根节点不与父节点之间的边被选中时，所选边的权值之和。

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

Python3JavaC++Go

class Solution:
    def maxScore(self, edges: List[List[int]]) -> int:
        def dfs(i):
            a = b = t = 0
            for j, w in g[i]:
                x, y = dfs(j)
                a += y
                b += y
                t = max(t, x - y + w)
            b += t
            return a, b

        g = defaultdict(list)
        for i, (p, w) in enumerate(edges[1:], 1):
            g[p].append((i, w))
        return dfs(0)[1]

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

    public long maxScore(int[][] edges) {
        int n = edges.length;
        g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (int i = 1; i < n; ++i) {
            int p = edges[i][0], w = edges[i][1];
            g[p].add(new int[] {i, w});
        }
        return dfs(0)[1];
    }

    private long[] dfs(int i) {
        long a = 0, b = 0, t = 0;
        for (int[] nxt : g[i]) {
            int j = nxt[0], w = nxt[1];
            long[] s = dfs(j);
            a += s[1];
            b += s[1];
            t = Math.max(t, s[0] - s[1] + w);
        }
        b += t;
        return new long[] {a, b};
    }
}

class Solution {
public:
    long long maxScore(vector<vector<int>>& edges) {
        int n = edges.size();
        vector<vector<pair<int, int>>> g(n);
        for (int i = 1; i < n; ++i) {
            int p = edges[i][0], w = edges[i][1];
            g[p].emplace_back(i, w);
        }
        using ll = long long;
        using pll = pair<ll, ll>;
        function<pll(int)> dfs = [&](int i) -> pll {
            ll a = 0, b = 0, t = 0;
            for (auto& [j, w] : g[i]) {
                auto [x, y] = dfs(j);
                a += y;
                b += y;
                t = max(t, x - y + w);
            }
            b += t;
            return make_pair(a, b);
        };
        return dfs(0).second;
    }
};

func maxScore(edges [][]int) int64 {
    n := len(edges)
    g := make([][][2]int, n)
    for i := 1; i < n; i++ {
        p, w := edges[i][0], edges[i][1]
        g[p] = append(g[p], [2]int{i, w})
    }
    var dfs func(int) [2]int
    dfs = func(i int) [2]int {
        var a, b, t int
        for _, e := range g[i] {
            j, w := e[0], e[1]
            s := dfs(j)
            a += s[1]
            b += s[1]
            t = max(t, s[0]-s[1]+w)
        }
        b += t
        return [2]int{a, b}
    }
    return int64(dfs(0)[1])
}

方法二：树形 DP + 显式栈

思考

在树中选一组不相邻的边，使权和最大。 \(n \le 10^5\)，从根递归决定每条边选或不选，链上的调用深度就是 \(n\)，会超出递归栈。

每个节点只要两个量：连向父节点的边被选中时的最大权和 \(a\)，以及这条边不选时的最大权和 \(b\)。父边被选中时，通向孩子的边都不能选， \(a\) 是各孩子 \(b\) 的和；父边不选时，最多再选一条通向孩子的边，增益是该孩子的 \(a\) 减去 \(b\) 再加上边权，取增益最大的一条。

显式栈按 \((节点, 状态)\) 做后序。进入时先压出栈标记再压孩子，离开时写下 \(a\) 和 \(b\)。根没有父边，答案是根的 \(b\)。

对每个节点 \(i\) 计算两个值 \((a, b)\)。 \(a\) 表示节点 \(i\) 与其父节点之间的边被选中时，子树里所选边的权值之和； \(b\) 表示这条边不被选中时的权值之和。用显式栈从根做后序，离开节点时再写这两个值。

对于当前节点 \(i\)：

如果 \(i\) 与父节点的边被选择，则它与子节点的所有边都不能被选择，那么 \(a\) 就是其所有子节点的 \(b\) 值之和；
如果 \(i\) 与父节点的边没被选择，那么可以选择它与子节点的最多一条边，那么 \(b\) 就是选中的那个子节点的 \(a\) 与其余子节点的 \(b\) 之和，再加上对应边的权值。增益为负时这条边不选。

根没有父边，答案是根的 \(b\)。