이진 트리(binary tree)가 하나 주어져 있다고 가정해 보겠습니다. 이때 구해야 할 것은 기존 트리와 구조는 같지만, 모든 노드의 값을 '자기 자신의 값 + 왼쪽 서브트리의 합 + 오른쪽 서브트리의 합'으로 대체한 새로운 트리입니다.
예를 들어 입력 트리가 다음과 같다면,
출력 결과는 아래와 같습니다.
문제 해결 접근 방법
이 문제는 재귀적인 후위 순회(post-order traversal)를 활용하면 간단하게 해결할 수 있습니다. 자식 노드들의 합계를 먼저 계산한 뒤 부모 노드의 값을 갱신하는 방식입니다. 알고리즘의 동작 순서는 다음과 같습니다.
tree_sum()함수를 정의합니다. 이 함수는 트리의 루트 노드를 인자로 받습니다.루트가 null(None)이라면 0을 반환합니다.
루트의 데이터 값을 다음과 같이 갱신합니다.
root.data = tree_sum(root.left) + tree_sum(root.right) + root.data갱신된 루트의 데이터를 반환합니다.
메인 메서드에서는
tree_sum(root)를 호출한 뒤root를 반환합니다.
구현 예제
아래 코드를 통해 더 잘 이해해 보겠습니다.
class TreeNode: def __init__(self, data, left=None, right=None): self.data = data self.left = left self.right = right def inorder(root): if root: inorder(root.left) print(root.data, end=', ') inorder(root.right) class Solution: def solve(self, root): def tree_sum(root: TreeNode): if root is None: return 0 root.data = tree_sum(root.left) + tree_sum(root.right) + root.data return root.data tree_sum(root) return root ob = Solution() root = TreeNode(2) root.left = TreeNode(3) root.right = TreeNode(4) root.left.left = TreeNode(9) root.left.right = TreeNode(7) ob.solve(root) inorder(root)
입력
root = TreeNode(2) root.left = TreeNode(3) root.right = TreeNode(4) root.left.left = TreeNode(9) root.left.right = TreeNode(7)
출력
9, 19, 7, 25, 4,
결과 설명
중위 순회(inorder) 순서로 살펴보면 각 노드의 값이 다음과 같이 변경됩니다.
값이 9였던 노드 → 리프 노드이므로 그대로 9
값이 3이었던 노드 → 3 + 9(왼쪽) + 7(오른쪽) = 19
값이 7이었던 노드 → 리프 노드이므로 그대로 7
값이 2였던 루트 노드 → 2 + 19(왼쪽 서브트리 전체 합) + 4(오른쪽) = 25
값이 4였던 노드 → 리프 노드이므로 그대로 4
이처럼 재귀 호출이 가장 깊은 곳(리프 노드)부터 값을 계산해 올라오면서 부모 노드에 누적되는 구조이기 때문에, 한 번의 트리 순회만으로 모든 노드를 O(n) 시간 복잡도 안에 갱신할 수 있습니다.