Algorithms training - DP on trees. 算法训练(二) 树上DP.

咩系树上DP? 其实都没咩可解释嘅. 即系在树上嘅每个节点表示状态, 然后推导状态转移方程.
正经D讲, 套路就是先”递”后”归”, 在”归”的过程中进行状态转移. 所以一般都是后序遍历的模板.

二叉树中的最大路径和

https://leetcode.cn/problems/binary-tree-maximum-path-sum/

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
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def maxPathSum(self, root: Optional[TreeNode]) -> int:
        ans = -inf

        def dfs(root: Optional[TreeNode]) -> int:
            nonlocal ans

            if root is None:
                return 0

            left = max(0, dfs(root.left))
            right = max(0, dfs(root.right))

            ans = max(ans, left + right + root.val)
            return max(left, right) + root.val
        
        dfs(root)

        return ans

二叉树的直径

https://leetcode.cn/problems/diameter-of-binary-tree/

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def diameterOfBinaryTree(self, root: Optional[TreeNode]) -> int:
        ans = 0

        def dfs(root: Optional[TreeNode]) -> int:
            nonlocal ans
            if root is None:
                return 0

            left = dfs(root.left)
            right = dfs(root.right)

            ans = max(ans, left + right + 1)
            return max(left, right) + 1
        dfs(root)
        return ans - 1 # 链的长度 = 节点数目 - 1

打家劫舍III

https://leetcode.cn/problems/house-robber-iii/

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution:
    def rob(self, root: Optional[TreeNode]) -> int:
        @cache
        def dfs(root: Optional[TreeNode]) -> (int, int):
            if root is None:
                return 0, 0

            l_rob, l_nrob = dfs(root.left)
            r_rob, r_nrob = dfs(root.right)

            rob = l_nrob + r_nrob + root.val # 偷
            nrob = max(l_rob, l_nrob) + max(r_rob, r_nrob) # 唔偷

            return rob, nrob
            

        return max(dfs(root))