跳转至

2920. 收集所有金币可获得的最大积分

来源第 369 场周赛 Q4难度困难分数2350

题目描述

有一棵由 n 个节点组成的无向树，以 0  为根节点，节点编号从 0 到 n - 1 。给你一个长度为 n - 1 的二维 整数 数组 edges ，其中 edges[i] = [ai, bi] 表示在树上的节点 ai 和 bi 之间存在一条边。另给你一个下标从 0 开始、长度为 n 的数组 coins 和一个整数 k ，其中 coins[i] 表示节点 i 处的金币数量。

从根节点开始，你必须收集所有金币。要想收集节点上的金币，必须先收集该节点的祖先节点上的金币。

节点 i 上的金币可以用下述方法之一进行收集：

  • 收集所有金币，得到共计 coins[i] - k 点积分。如果 coins[i] - k 是负数，你将会失去 abs(coins[i] - k) 点积分。
  • 收集所有金币，得到共计 floor(coins[i] / 2) 点积分。如果采用这种方法，节点 i 子树中所有节点 j 的金币数 coins[j] 将会减少至 floor(coins[j] / 2) 。

返回收集 所有 树节点的金币之后可以获得的最大积分。

 

示例 1：

输入：edges = [[0,1],[1,2],[2,3]], coins = [10,10,3,3], k = 5
输出：11                        
解释：
使用第一种方法收集节点 0 上的所有金币。总积分 = 10 - 5 = 5 。
使用第一种方法收集节点 1 上的所有金币。总积分 = 5 + (10 - 5) = 10 。
使用第二种方法收集节点 2 上的所有金币。所以节点 3 上的金币将会变为 floor(3 / 2) = 1 ，总积分 = 10 + floor(3 / 2) = 11 。
使用第二种方法收集节点 3 上的所有金币。总积分 =  11 + floor(1 / 2) = 11.
可以证明收集所有节点上的金币能获得的最大积分是 11 。

示例 2：

输入：edges = [[0,1],[0,2]], coins = [8,4,4], k = 0
输出：16
解释：
使用第一种方法收集所有节点上的金币，因此，总积分 = (8 - 0) + (4 - 0) + (4 - 0) = 16 。

 

提示：

  • n == coins.length
  • 2 <= n <= 105
  • 0 <= coins[i] <= 104
  • edges.length == n - 1
  • 0 <= edges[i][0], edges[i][1] < n
  • 0 <= k <= 104

解法

方法一：记忆化搜索

思考

树上每个点可选「减 \(k\)」或「全体再右移一位」。后一种操作对整棵尚未处理的子树生效，且金币 \(\le 10^4\)，右移约 \(14\) 次后归零，故位移次数是很小的额外状态。

\(dfs(i,fa,j)\) 表示当前位移为 \(j\) 时在 \(i\) 的最优得分：方案甲取 \((coins[i] \gg j)-k\) 且子节点仍用 \(j\)；方案乙取 \(coins[i] \gg (j+1)\) 且子节点用 \(j+1\)（以 \(j<14\) 为限）。记忆化后从根出发即可。

我们先根据题目给定的边构建图 \(g\)，其中 \(g[i]\) 表示节点 \(i\) 的所有邻接节点。然后我们可以使用记忆化搜索的方法求解本题。

我们设计一个函数 \(dfs(i, fa, j)\)，表示当前节点为 \(i\)，父节点为 \(fa\)，当前节点的金币数需要右移 \(j\) 位，所能获得的最大积分。

函数 \(dfs(i, fa, j)\) 的执行过程如下：

如果我们使用第一种方法收集当前节点的金币，那么当前节点的积分为 \((coins[i] >> j) - k\)，然后我们遍历当前节点的所有邻接节点 \(c\)，如果 \(c\) 不等于 \(fa\)，那么我们将 \(dfs(c, i, j)\) 的结果累加到当前节点的积分中。

如果我们使用第二种方法收集当前节点的金币，那么当前节点的积分为 \(coins[i] >> (j + 1)\)，然后我们遍历当前节点的所有邻接节点 \(c\)，如果 \(c\) 不等于 \(fa\)，那么我们将 \(dfs(c, i, j + 1)\) 的结果累加到当前节点的积分中。注意，由于题目中给定的 \(coins[i]\) 最大值为 \(10^4\)，因此我们最多只能右移 \(14\) 位，就使得 \(coins[i] >> (j + 1)\) 的值为 \(0\)。

最后，我们返回当前节点使用两种方法中能获得的最大积分。

为了避免重复计算，我们使用记忆化搜索的方法，将 \(dfs(i, fa, j)\) 的结果存储到 \(f[i][j]\) 中，其中 \(f[i][j]\) 表示当前节点为 \(i\)，父节点为 \(fa\)，当前节点的金币数需要右移 \(j\) 位，所能获得的最大 时间复杂度 \(O(n \times \log M)\)，空间复杂度 \(O(n \times \log M)\)。其中 \(M\) 表示 \(coins[i]\) 的最大值。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:
        @cache
        def dfs(i: int, fa: int, j: int) -> int:
            a = (coins[i] >> j) - k
            b = coins[i] >> (j + 1)
            for c in g[i]:
                if c != fa:
                    a += dfs(c, i, j)
                    if j < 14:
                        b += dfs(c, i, j + 1)
            return max(a, b)

        n = len(coins)
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = dfs(0, -1, 0)
        dfs.cache_clear()
        return ans

 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
class Solution {
    private int k;
    private int[] coins;
    private Integer[][] f;
    private List<Integer>[] g;

    public int maximumPoints(int[][] edges, int[] coins, int k) {
        this.k = k;
        this.coins = coins;
        int n = coins.length;
        f = new Integer[n][15];
        g = new List[n];
        Arrays.setAll(g, i -> 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, 0);
    }

    private int dfs(int i, int fa, int j) {
        if (f[i][j] != null) {
            return f[i][j];
        }
        int a = (coins[i] >> j) - k;
        int b = coins[i] >> (j + 1);
        for (int c : g[i]) {
            if (c != fa) {
                a += dfs(c, i, j);
                if (j < 14) {
                    b += dfs(c, i, j + 1);
                }
            }
        }
        return f[i][j] = Math.max(a, 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
28
29
30
31
class Solution {
public:
    int maximumPoints(vector<vector<int>>& edges, vector<int>& coins, int k) {
        int n = coins.size();
        int f[n][15];
        memset(f, -1, sizeof(f));
        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);
        }
        auto dfs = [&](this auto&& dfs, int i, int fa, int j) -> int {
            if (f[i][j] != -1) {
                return f[i][j];
            }
            int a = (coins[i] >> j) - k;
            int b = coins[i] >> (j + 1);
            for (int c : g[i]) {
                if (c != fa) {
                    a += dfs(c, i, j);
                    if (j < 14) {
                        b += dfs(c, i, j + 1);
                    }
                }
            }
            return f[i][j] = max(a, b);
        };
        return dfs(0, -1, 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
func maximumPoints(edges [][]int, coins []int, k int) int {
    n := len(coins)
    f := make([][]int, n)
    for i := range f {
        f[i] = make([]int, 15)
        for j := range f[i] {
            f[i][j] = -1
        }
    }
    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)
    }
    var dfs func(int, int, int) int
    dfs = func(i, fa, j int) int {
        if f[i][j] != -1 {
            return f[i][j]
        }
        a := (coins[i] >> j) - k
        b := coins[i] >> (j + 1)
        for _, c := range g[i] {
            if c != fa {
                a += dfs(c, i, j)
                if j < 14 {
                    b += dfs(c, i, j+1)
                }
            }
        }
        f[i][j] = max(a, b)
        return f[i][j]
    }
    return dfs(0, -1, 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
function maximumPoints(edges: number[][], coins: number[], k: number): number {
    const n = coins.length;
    const f: number[][] = Array.from({ length: n }, () => Array(15).fill(-1));
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const dfs = (i: number, fa: number, j: number): number => {
        if (f[i][j] !== -1) {
            return f[i][j];
        }
        let a = (coins[i] >> j) - k;
        let b = coins[i] >> (j + 1);
        for (const c of g[i]) {
            if (c !== fa) {
                a += dfs(c, i, j);
                if (j < 14) {
                    b += dfs(c, i, j + 1);
                }
            }
        }
        return (f[i][j] = Math.max(a, b));
    };
    return dfs(0, -1, 0);
}

方法二：树形 DP + 显式栈

思考

树上每个点可选「减 \(k\)」或「再右移一位」。后一种操作对尚未处理的子树生效，金币不超过 \(10^4\)，大约右移 \(14\) 次后归零。 \(n \le 10^5\)，从根按同一位移递归下去，链上的调用深度就是 \(n\)，会超出递归栈。

节点 \(i\) 在位移 \(j\) 下的最优得分只依赖子节点在位移 \(j\) 和 \(j + 1\) 下的得分。方案甲取 \((coins[i] \gg j) - k\)，子节点仍用 \(j\)；方案乙取 \(coins[i] \gg (j + 1)\)，子节点用 \(j + 1\)，并且 \(j\) 不超过 \(14\)。

显式栈按 \((节点, 父节点, 状态)\) 做后序。孩子离开之后，再为当前点填上 \(0\) 到 \(14\) 的整张表。根在位移 \(0\) 下的值就是答案。

我们先根据题目给定的边构建图 \(g\)，其中 \(g[i]\) 表示节点 \(i\) 的所有邻接节点。

\(f[i][j]\) 表示节点 \(i\) 的金币已经右移 \(j\) 位时，子树能获得的最大积分。用显式栈从根做后序，孩子的表填完以后再填当前点。对每个位移 \(j\)（ \(0 \le j \le 14\)）：

如果我们使用第一种方法收集当前节点的金币，那么当前节点的积分为 \((coins[i] >> j) - k\)，再累加每个子节点 \(c\) 的 \(f[c][j]\)。

如果我们使用第二种方法收集当前节点的金币，那么当前节点的积分为 \(coins[i] >> (j + 1)\)。当 \(j < 14\) 时，再累加每个子节点的 \(f[c][j + 1]\)。金币不超过 \(10^4\)，右移 \(14\) 位后剩余值已经是 \(0\)。

\(f[i][j]\) 取这两种方法的较大值。答案是 \(f[0][0]\)。

时间复杂度 \(O(n \times \log M)\)，空间复杂度 \(O(n \times \log M)\)。其中 \(M\) 表示 \(coins[i]\) 的最大值。

 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
class Solution:
    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:
        n = len(coins)
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        f = [[0] * 15 for _ in range(n)]
        stk = [(0, -1, 0)]
        while stk:
            i, fa, state = stk.pop()
            if state == 0:
                stk.append((i, fa, 1))
                for c in g[i]:
                    if c != fa:
                        stk.append((c, i, 0))
            else:
                for j in range(15):
                    a = (coins[i] >> j) - k
                    b = coins[i] >> (j + 1)
                    for c in g[i]:
                        if c != fa:
                            a += f[c][j]
                            if j < 14:
                                b += f[c][j + 1]
                    f[i][j] = max(a, b)
        return f[0][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
class Solution {
    public int maximumPoints(int[][] edges, int[] coins, int k) {
        int n = coins.length;
        List<Integer>[] g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        int[][] f = new int[n][15];
        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 c : g[i]) {
                    if (c != fa) {
                        stk.push(new int[] {c, i, 0});
                    }
                }
            } else {
                for (int j = 0; j < 15; ++j) {
                    int a = (coins[i] >> j) - k;
                    int b = coins[i] >> (j + 1);
                    for (int c : g[i]) {
                        if (c != fa) {
                            a += f[c][j];
                            if (j < 14) {
                                b += f[c][j + 1];
                            }
                        }
                    }
                    f[i][j] = Math.max(a, b);
                }
            }
        }
        return f[0][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
class Solution {
public:
    int maximumPoints(vector<vector<int>>& edges, vector<int>& coins, int k) {
        int n = coins.size();
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        vector<array<int, 15>> f(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 c : g[i]) {
                    if (c != fa) {
                        stk.push_back({c, i, 0});
                    }
                }
            } else {
                for (int j = 0; j < 15; ++j) {
                    int a = (coins[i] >> j) - k;
                    int b = coins[i] >> (j + 1);
                    for (int c : g[i]) {
                        if (c != fa) {
                            a += f[c][j];
                            if (j < 14) {
                                b += f[c][j + 1];
                            }
                        }
                    }
                    f[i][j] = max(a, b);
                }
            }
        }
        return f[0][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
func maximumPoints(edges [][]int, coins []int, k int) int {
    n := len(coins)
    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)
    }
    f := make([][15]int, n)
    stk := [][3]int{{0, -1, 0}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        i, fa, state := cur[0], cur[1], cur[2]
        if state == 0 {
            stk = append(stk, [3]int{i, fa, 1})
            for _, c := range g[i] {
                if c != fa {
                    stk = append(stk, [3]int{c, i, 0})
                }
            }
        } else {
            for j := 0; j < 15; j++ {
                a := (coins[i] >> j) - k
                b := coins[i] >> (j + 1)
                for _, c := range g[i] {
                    if c != fa {
                        a += f[c][j]
                        if j < 14 {
                            b += f[c][j+1]
                        }
                    }
                }
                f[i][j] = max(a, b)
            }
        }
    }
    return f[0][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
function maximumPoints(edges: number[][], coins: number[], k: number): number {
    const n = coins.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const f: number[][] = Array.from({ length: n }, () => Array(15).fill(0));
    const stk: [number, number, number][] = [[0, -1, 0]];
    while (stk.length) {
        const [i, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([i, fa, 1]);
            for (const c of g[i]) {
                if (c !== fa) {
                    stk.push([c, i, 0]);
                }
            }
        } else {
            for (let j = 0; j < 15; ++j) {
                let a = (coins[i] >> j) - k;
                let b = coins[i] >> (j + 1);
                for (const c of g[i]) {
                    if (c !== fa) {
                        a += f[c][j];
                        if (j < 14) {
                            b += f[c][j + 1];
                        }
                    }
                }
                f[i][j] = Math.max(a, b);
            }
        }
    }
    return f[0][0];
}

评论