337. House Robber III

House Robber III - LeetCode

1. BFS를 이용한 Level 별 탐색(실패)

from collections import defaultdict, deque


class Solution:
    def rob(self, root: TreeNode) -> int:
        dict1 = defaultdict(int)
        if not root:
            return 0
        q = deque([(root, 0)])
        while q:
            cur_node, cur_level = q.popleft()
            dict1[cur_level] += cur_node.val
            if cur_node.right:
                q.append((cur_node.right, cur_level + 1))
            if cur_node.left:
                q.append((cur_node.left, cur_level + 1))
        answer1, answer2 = 0, 0
        for level in dict1:
            if level % 2 == 0:
                answer1 += dict1[level]
            else:
                answer2 += dict1[level]
        return max(answer1, answer2)

트리를 BFS로 탐색하면서 동일 레벨의 노드 값을 저장한다.

그 후 저장된 레벨을 홀수, 짝수로 묶고 그중 큰 값을 반환한다.

이 로직의 반례는 아래와 같은 그래프이다.

 

위에서부터 1,2,3,4번 노드라고 했을 때 홀수, 짝수 레벨별로 생각하는 로직은 (1,3), (2,4)로 생각하지만 (1,4), (2,3)은 고려하지 않는다.

2. Tree DP

class Solution:
    def rob(self, root: TreeNode) -> int:
        def dfs(cur):
            if not cur:
                return [0, 0]
            left = dfs(cur.left)
            right = dfs(cur.right)

            return [max(left) + max(right), left[0] + right[0] + cur.val]

        answer = dfs(root)
        return max(answer)

Tree DP는 Tree와 DP를 조합해서 문제를 푸는 방법이다.

주로 특정 서브 트리의 결과를 이용해야 할 때 사용하는 것 같다.

 

생각해보면 이렇다.

dp[node][0] = max(dp[node.left][0], dp[node.left][1]) + max(dp[node.right][0], dp[node.right][1])

현재 node를 사용하지 않았을 때(0)의 최댓값

=

현재 node의 왼쪽 노드를 사용하지 않았을 때의 왼쪽 서브 트리의 값과

현재 node의 왼쪽 노드를 사용했을 때의 왼쪽 서브 트리의 값 중 최댓값 

+

현재 node의 오른쪽 노드를 사용하지 않았을 때의 오른쪽 서브 트리의 값과

현재 node의 오른쪽 노드를 사용했을 때의 오른쪽 서브 트리의 값 중 최댓값 

 

현재 노드를 사용하지 않는다는 것은 왼쪽, 오른쪽 서브 트리의 루트 노드(즉 현재 노드와 바로 연결되어있는 노드)를 사용하던 말던 상관없다.

 

dp[node][1] = dp[node.left][0] + dp[node.right][0] + node.val

현재 node를 사용할 때(1)의 최댓값

=

현재 node의 왼쪽 노드를 사용하지 않았을 때의 왼쪽 서브 트리의 값

+

현재 node의 오른쪽 노드를 사용하지 않았을 때의 오른쪽 서브 트리의 값

+

현재 node의 값

 

현재 노드를 사용한다는 것은 왼쪽, 오른쪽 서브 트리의 루트 노드를 사용하지 않은 값들의 합 + 현재 노드의 값이 최댓값이 된다.

 

결국 위 점화식을 최하단 자식 노드부터 올라가면서 DP 테이블을 완성해나가면

[루트 노드를 사용하지 않았을 때의 최댓값, 루트 노드를 사용했을 때의 최댓값]을 얻을 수 있고 이 중 큰 값을 반환하면 된다.

 

 

반례를 보고 어떻게 풀어야 하는지 알아낸 것이 좀 아쉽지만 저 코드를 내가 직접 생각해서 빠른 시간에 작성해서 기분이 좋았다. 꽤 괜찮은 코드를 오랜만에 쓴 것 같다.