跳转至

2925. 在树上执行操作以后得到的最大分数

来源第 370 场周赛 Q3难度中等分数1939

题目描述

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

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

一开始你的分数为 0。在一次操作中，你选择一个节点 i，将 values[i] 的值 加 到你的得分上，然后将 values[i] 设置 为 0。这三个步骤作为一个单一操作同时发生。

如果从根节点出发，到 任意 叶子节点经过的路径上的节点值之和都 不等于 0 ，那么我们称这棵树是 健康的 。

你可以对这棵树执行任意次操作，但要求执行完所有操作以后树是 健康的 ，请你返回你可以获得的 最大分数 。

 

示例 1：

输入：edges = [[0,1],[0,2],[0,3],[2,4],[4,5]], values = [5,2,5,2,1,1]
输出：11
解释：我们对节点 1、2、3、4 和 5 进行操作，因此值变为 [5,0,0,0,0,0]。叶子节点是节点 1、3 和 5。
- 从 0 到 1 的路径上值的总和等于 5。
- 从 0 到 3 的路径上值的总和等于 5。
- 从 0 到 5 的路径上值的总和等于 5。
每个叶子节点都有非零的路径和，所以树是健康的。得分是所选节点的原始值的总和：2 + 5 + 2 + 1 + 1 = 11。
可以证明 11 是通过在树上执行任何数量的操作所能获得的最大分数。

示例 2：

输入：edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], values = [20,10,9,7,4,3,5]
输出：40
解释：我们对节点 0、2、3 和 4 进行操作，因此值变为 [0,10,0,0,0,3,5]。叶子节点是节点 3、4、5 和 6。
- 从 0 到 3 的路径上值的总和等于 10。
- 从 0 到 4 的路径上值的总和等于 10。
- 从 0 到 5 的路径上值的总和等于 3。
- 从 0 到 6 的路径上值的总和等于 5。
每个叶子节点都有非零的路径和，所以树是健康的。得分是所选节点原始值的总和：20 + 9 + 7 + 4 = 40。
可以证明 40 是通过在树上执行任何数量的操作所能获得的最高分数。

 

提示：

  • 2 <= n <= 2 * 104
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • values.length == n
  • 1 <= values[i] <= 109
  • 输入保证 edges 构成一棵合法的树。

解法

方法一：树形 DP

思考

树上每条根到叶的路径必须留下至少一个未选节点，其余节点值可计入得分。若直接枚举每个点选或不选，需保证所有路径约束，状态过多。树形 DP 把约束收束到子树：要么不选根，则子树可全选；要么选根，则每个子树内部仍须各自留下一个未选点。

\(dfs\) 返回子树总和以及「满足路径约束的最大选取」。叶子只能不选自身才能留下未选点，故第二维为 \(0\)。非叶取 \(\max(values[i]+b, a)\)，其中 \(a\) 为子树全和、 \(b\) 为子树均已合法。答案是根的第二维。

题目实际上是让我们从树的所有节点中选出一些节点，使得这些节点的值之和最大，并且每条从根节点到叶子节点的路径上都有一个点没有被选中。

我们可以使用树形 DP 的方法解决这个问题。

我们设计一个函数 \(dfs(i, fa)\)，其中 \(i\) 表示当前以节点 \(i\) 作为子树的根节点，且 \(fa\) 表示 \(i\) 的父节点，函数返回一个长度为 \(2\) 的数组，其中 \([0]\) 表示该子树中所有节点的值之和，而 \([1]\) 表示该子树满足每条路径上都有一个点没有被选中的最大值。

其中 \([0]\) 的值可以直接通过 DFS 累加每个节点的值得到，而 \([1]\) 的值，则需要考虑两种情况，即节点 \(i\) 是否被选中。如果被选中，那么节点 \(i\) 的每个子树得必须满足每条路径上都有一个点没有被选中；如果没有被选中，那么节点 \(i\) 的每个子树可以选取所有节点。我们取这两种情况中的最大值即可。

需要注意的是，叶子节点的 \([1]\) 的值为 \(0\)，因为叶子节点没有子树，所以不需要考虑每条路径上都有一个点没有被选中的情况。

答案为 \(dfs(0, -1)[1]\)。

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
    def maximumScoreAfterOperations(
        self, edges: List[List[int]], values: List[int]
    ) -> int:
        def dfs(i: int, fa: int = -1) -> (int, int):
            a = b = 0
            leaf = True
            for j in g[i]:
                if j != fa:
                    leaf = False
                    aa, bb = dfs(j, i)
                    a += aa
                    b += bb
            if leaf:
                return values[i], 0
            return values[i] + a, max(values[i] + b, a)

        g = [[] for _ in range(len(values))]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        return dfs(0)[1]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
    private List<Integer>[] g;
    private int[] values;

    public long maximumScoreAfterOperations(int[][] edges, int[] values) {
        int n = values.length;
        g = new List[n];
        this.values = values;
        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);
        }
        return dfs(0, -1)[1];
    }

    private long[] dfs(int i, int fa) {
        long a = 0, b = 0;
        boolean leaf = true;
        for (int j : g[i]) {
            if (j != fa) {
                leaf = false;
                var t = dfs(j, i);
                a += t[0];
                b += t[1];
            }
        }
        if (leaf) {
            return new long[] {values[i], 0};
        }
        return new long[] {values[i] + a, Math.max(values[i] + b, a)};
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution {
public:
    long long maximumScoreAfterOperations(vector<vector<int>>& edges, vector<int>& values) {
        int n = values.size();
        vector<int> g[n];
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].emplace_back(b);
            g[b].emplace_back(a);
        }
        using ll = long long;
        function<pair<ll, ll>(int, int)> dfs = [&](int i, int fa) -> pair<ll, ll> {
            ll a = 0, b = 0;
            bool leaf = true;
            for (int j : g[i]) {
                if (j != fa) {
                    auto [aa, bb] = dfs(j, i);
                    a += aa;
                    b += bb;
                    leaf = false;
                }
            }
            if (leaf) {
                return {values[i], 0LL};
            }
            return {values[i] + a, max(values[i] + b, a)};
        };
        auto [_, b] = dfs(0, -1);
        return b;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
func maximumScoreAfterOperations(edges [][]int, values []int) int64 {
    g := make([][]int, len(values))
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    var dfs func(int, int) (int64, int64)
    dfs = func(i, fa int) (int64, int64) {
        a, b := int64(0), int64(0)
        leaf := true
        for _, j := range g[i] {
            if j != fa {
                leaf = false
                aa, bb := dfs(j, i)
                a += aa
                b += bb
            }
        }
        if leaf {
            return int64(values[i]), int64(0)
        }
        return int64(values[i]) + a, max(int64(values[i])+b, a)
    }
    _, b := dfs(0, -1)
    return b
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
function maximumScoreAfterOperations(edges: number[][], values: number[]): number {
    const g: number[][] = Array.from({ length: values.length }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const dfs = (i: number, fa: number): [number, number] => {
        let [a, b] = [0, 0];
        let leaf = true;
        for (const j of g[i]) {
            if (j !== fa) {
                const [aa, bb] = dfs(j, i);
                a += aa;
                b += bb;
                leaf = false;
            }
        }
        if (leaf) {
            return [values[i], 0];
        }
        return [values[i] + a, Math.max(values[i] + b, a)];
    };
    return dfs(0, -1)[1];
}

方法二：树形 DP + 显式栈

思考

每条从根到叶子的路径都必须留下至少一个未选节点，其余节点的值可以计入得分。若对每个节点分别枚举选或不选，还要同时满足全部路径约束，状态规模随路径数增长，在 \(n \le 2 \times 10^4\) 时无法承受。把决策收束到子树之后，当前节点不选则各子树可以整棵取走，当前节点选中则每个子树内部仍须各自合法；这两种选择只依赖子树全和与合法得分，与子节点的处理顺序无关。沿一条链递归填写这对数值时，调用深度等于节点数，会超出 Python 的递归上限。因此栈中保存 \((节点, 父节点, 状态)\)：状态 \(0\) 先压入退出标记再压入子节点，状态 \(1\) 在子结果就绪后写入当前节点。叶子的第二维为 \(0\)，非叶取 \(\max(values[i]+b, a)\)，答案是根的第二维。

题目要求从树中选出若干节点，使它们的值之和最大，并且每条从根到叶子的路径上都至少有一个节点没有被选中。

我们用树形 DP 记录每个子树的两个量，并用显式栈按后序填写。对节点 \(i\)，第一维是子树中所有节点的值之和，第二维是该子树在满足路径约束时能取得的最大得分。叶子没有子节点，只能不选自身才能留下未选点，因此第一维为 \(values[i]\)，第二维为 \(0\)。非叶节点有两种取法：不选 \(i\) 时可以取走每个子树的全部节点，得分为各子树之和 \(a\)；选中 \(i\) 时每个子树内部仍须各自合法，得分为 \(values[i]\) 加上各子树的合法得分 \(b\)。取二者的较大值。

栈中的每个元素是 \((i, fa, state)\)。 \(state = 0\) 时先压入 \((i, fa, 1)\)，再把除父节点以外的邻接点以状态 \(0\) 压入，这样子节点先完成。 \(state = 1\) 时读取已经写好的子节点结果，按上面的规则写入节点 \(i\)。子节点入栈的先后不影响答案。

答案为根节点的第二维。

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution:
    def maximumScoreAfterOperations(
        self, edges: List[List[int]], values: List[int]
    ) -> int:
        n = len(values)
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        sub = [(0, 0)] * n
        stk = [(0, -1, 0)]
        while stk:
            i, fa, state = stk.pop()
            if state == 0:
                stk.append((i, fa, 1))
                for j in g[i]:
                    if j != fa:
                        stk.append((j, i, 0))
            else:
                a = b = 0
                leaf = True
                for j in g[i]:
                    if j != fa:
                        leaf = False
                        aa, bb = sub[j]
                        a += aa
                        b += bb
                if leaf:
                    sub[i] = (values[i], 0)
                else:
                    sub[i] = (values[i] + a, max(values[i] + b, a))
        return sub[0][1]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class Solution {
    public long maximumScoreAfterOperations(int[][] edges, int[] values) {
        int n = values.length;
        List<Integer>[] g = new List[n];
        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);
        }
        long[] sum = new long[n];
        long[] best = new long[n];
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push(new int[] {i, fa, 1});
                for (int j : g[i]) {
                    if (j != fa) {
                        stk.push(new int[] {j, i, 0});
                    }
                }
            } else {
                long a = 0, b = 0;
                boolean leaf = true;
                for (int j : g[i]) {
                    if (j != fa) {
                        leaf = false;
                        a += sum[j];
                        b += best[j];
                    }
                }
                if (leaf) {
                    sum[i] = values[i];
                    best[i] = 0;
                } else {
                    sum[i] = values[i] + a;
                    best[i] = Math.max(values[i] + b, a);
                }
            }
        }
        return best[0];
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
public:
    long long maximumScoreAfterOperations(vector<vector<int>>& edges, vector<int>& values) {
        int n = values.size();
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].emplace_back(b);
            g[b].emplace_back(a);
        }
        using ll = long long;
        vector<pair<ll, ll>> sub(n);
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto [i, fa, state] = stk.back();
            stk.pop_back();
            if (state == 0) {
                stk.push_back({i, fa, 1});
                for (int j : g[i]) {
                    if (j != fa) {
                        stk.push_back({j, i, 0});
                    }
                }
            } else {
                ll a = 0, b = 0;
                bool leaf = true;
                for (int j : g[i]) {
                    if (j != fa) {
                        leaf = false;
                        a += sub[j].first;
                        b += sub[j].second;
                    }
                }
                if (leaf) {
                    sub[i] = {values[i], 0};
                } else {
                    sub[i] = {values[i] + a, max(values[i] + b, a)};
                }
            }
        }
        return sub[0].second;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
func maximumScoreAfterOperations(edges [][]int, values []int) int64 {
    n := len(values)
    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)
    }
    type pair struct{ sum, best int64 }
    sub := make([]pair, n)
    type frame struct{ i, fa, state int }
    stk := []frame{{0, -1, 0}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        i, fa, state := cur.i, cur.fa, cur.state
        if state == 0 {
            stk = append(stk, frame{i, fa, 1})
            for _, j := range g[i] {
                if j != fa {
                    stk = append(stk, frame{j, i, 0})
                }
            }
        } else {
            var a, b int64
            leaf := true
            for _, j := range g[i] {
                if j != fa {
                    leaf = false
                    a += sub[j].sum
                    b += sub[j].best
                }
            }
            if leaf {
                sub[i] = pair{int64(values[i]), 0}
            } else {
                sub[i] = pair{int64(values[i]) + a, max(int64(values[i])+b, a)}
            }
        }
    }
    return sub[0].best
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
function maximumScoreAfterOperations(edges: number[][], values: number[]): number {
    const n = values.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const sum = Array(n).fill(0);
    const best = Array(n).fill(0);
    const stk: number[][] = [[0, -1, 0]];
    while (stk.length) {
        const [i, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([i, fa, 1]);
            for (const j of g[i]) {
                if (j !== fa) {
                    stk.push([j, i, 0]);
                }
            }
        } else {
            let a = 0;
            let b = 0;
            let leaf = true;
            for (const j of g[i]) {
                if (j !== fa) {
                    leaf = false;
                    a += sum[j];
                    b += best[j];
                }
            }
            if (leaf) {
                sum[i] = values[i];
                best[i] = 0;
            } else {
                sum[i] = values[i] + a;
                best[i] = Math.max(values[i] + b, a);
            }
        }
    }
    return best[0];
}

评论