摘要: 等效性剪枝优化树形 DP

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

各位好，今天我们来看一个难题拆解，算法的主框架为树形 DP，第一个难点是发现子树作为阶段需要再加上附加信息才能具备最优子结构，另一个难点是状态空间中存在大量可以被等效掉的状态，这些状态的结果是已知的，无需去计算而代入已知的结果即可，DP 推导过程中的这部分状态空间自然就省掉了，从剪枝的类型来看，这是一种等效性剪枝。

题目

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

有一棵由 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
2
3
4
5
6

n == coins.length
2 <= n <= 1e5
0 <= coins[i] <= 1e4
edges.length == n - 1
0 <= edges[i][0], edges[i][1] < n
0 <= k <= 1e4

示例 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 。

题解

算法: 树形DP

考虑树中的节点 $u$，假设 $u$ 有 $s$ 个子节点，分别记为 $v_{1}, \cdots, v_{s}$。如果忽略 $u$ 的父节点 $f$ 带来的 影响，仅考虑 $u$ 这棵子树自身最多能获得多少积分，这就是一个子问题。

但仅考虑子树自身尚不能形成最优子结构，也就是原问题（即根节点对应的树可以获得的最大积分）不一定包含 $u$ 这棵子树视为孤立的树时的最大积分。

因为父节点有其它信息会对 $u$ 这棵子树上的结果有影响，这个信息就是 $u$ 的祖先链上共做了多少次操作 (2)，把这个信息加上，就可以形成最优子结构了。

综上，这是阶段（子树）和附加信息（祖先链上操作 (2) 的次数）共同形成的最优子结构，其状态表示为 $dp[u][j]$，也就时当 $u$ 的祖先链上的操作 (2) 次数为 $j$ 时，$u$ 对应的子树可以取得的最大积分。

状态这样表示的话，$dp[root][0]$ 为问题的答案，叶节点对应初始状态。状态转移方程如下：

其中 $g(u, j)$ 表示当祖先链上有 $j$ 次操作 (2) 时，$c[u]$ 会减少至多少，若 $c[u]$ 本身为 0，则 $g(u, j) = 0$。

代码 (Python)

注：DP 推导的无效状态太多，空间超限，需要对状态空间剪枝。

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

from functools import lru_cache

class Solution:
    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:
        n = len(coins)
        graph = [[] for _ in range(n)]
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        @lru_cache(maxsize=int(1e7))
        def g(u: int, j: int) -> int:
            if j == 0:
                return coins[u]
            return g(u, j - 1) // 2

        @lru_cache(maxsize=int(1e7))
        def solve(fa: int, u: int, j: int) -> int:
            ans1 = g(u, j) - k
            ans2 = g(u, j) // 2
            for v in graph[u]:
                if v == fa:
                    continue
                ans1 += solve(u, v, j)
                ans2 += solve(u, v, j + 1)
            return max(ans1, ans2)

        return solve(-1, 0, 0)

优化

如果只是上述的树形 DP 算法，本题还不至于很难。后面的两个优化是通过本题的关键。

优化1：移位操作

一次操作 (2) 要做一次 $x = \lfloor\frac{x}{2}\rfloor$，这相当于 x = (x >> 1)。

另一方面，移位操作可以叠加，也就是 x >> 1 >> 1 = x >> 2。于是 $j$ 次操作 (2) 后，自然有 x = (x >> j)。

这样，计算 $g(u, j)$ 就变得简单了，直接 g(u, j) = coins[u] >> j 即可。

优化2：剪枝优化 DP

如前所述，DP 推导的无效状态太多，空间超限，需要对状态空间剪枝。

由于树中各节点的值是有范围的，它们移位若干次后，就都变为零了，此时再进行操作 (2)，子树上得到的结果只会是 0，这样就不用去计算 2 对应的状态了，这部分状态空间就可以省了。

由数据范围给定的 0 <= coins[i] <= 1e4，其对应的二进制最多 14 位，也就是最多移位 14 次，就进入了上面说的操作 (2) 对应的结果直接为 0 的情况了。于是可以增加等效性剪枝：当 j > 14，跳过操作 (2) 的计算，直接将 0 作为对应的结果即可。

可以想象，如果没有这个剪枝，如果树的形状是高度非常高的那种，$j$ 这一维度的状态就会撑得很大，但如前所述，其中大于 14 的部分都是无效的。

代码 (Python)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

from functools import lru_cache

class Solution:
    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:
        n = len(coins)
        graph = [[] for _ in range(n)]
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        @lru_cache(maxsize=int(1e7))
        def solve(fa: int, u: int, j: int) -> int:
            ans1 = (coins[u] >> j) - k
            ans2 = coins[u] >> (j + 1)
            for v in graph[u]:
                if v == fa:
                    continue
                ans1 += solve(u, v, j)
                if j <= 14:
                    ans2 += solve(u, v, j + 1)
            return max(ans1, ans2)

        return solve(-1, 0, 0)