337. House Robber III

[zeikar]

Problem link

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

Problem Summary

이진 트리 상에서 연속하지 않은 노드의 합 중 최댓값을 구하는 문제.

Solution

딱 보니 간단한 DP이다.

현재 노드에서 훔치면 다음 자식 노드에서는 훔치면 안되고 현재 노드에서 훔치지 않았으면 자식 노드에서는 훔치거나 훔치지 않거나 2가지를 탐색하면 된다.
연속되지 않아야 하므로 이전의 상태가 필요하고 이도 같이 메모이제이션 해주면 된다. 파이썬은 간단하게 lru_cache가 있어 설정만 해주면 된다.
(파이썬으로 문제 풀이 익숙해지면 C++보다 편하긴 하다)

Source Code

from functools import lru_cache
from typing import Optional


# 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 rob(self, root: Optional[TreeNode]) -> int:
        return self.solve(root, False)

    @lru_cache(None)
    def solve(self, root: Optional[TreeNode], stole_before: bool) -> int:
        if not root:
            return 0
        if stole_before:
            return self.solve(root.left, False) + self.solve(root.right, False)
        else:
            return max(self.solve(root.left, True) + self.solve(root.right, True) + root.val,
                       self.solve(root.left, False) + self.solve(root.right, False))