파이썬으로 이진 트리의 모든 노드 값 합계 구하는 방법

값을 가진 이진 트리(binary tree)가 주어졌을 때, 트리에 포함된 모든 노드 값의 합계를 구해야 하는 경우가 있습니다.

예를 들어 다음과 같은 트리가 입력으로 주어지면

출력 결과는 14가 됩니다. 즉, 2 + 4 + 3 + 5 = 14입니다.

문제 해결 접근 방법

이 문제는 재귀(Recursion)를 활용하면 간단하게 해결할 수 있습니다. 트리 순회는 본질적으로 재귀적인 구조를 가지기 때문입니다. 해결 단계는 다음과 같습니다.

• 노드(node)를 인자로 받는 recurse() 함수를 정의합니다.
• val 변수에 현재 노드의 값을 저장합니다.
• 노드의 왼쪽 자식이 존재하면, val에 왼쪽 서브트리의 합계(recurse(left))를 더합니다.
• 노드의 오른쪽 자식이 존재하면, val에 오른쪽 서브트리의 합계(recurse(right))를 더합니다.
• val을 반환합니다.

메인 메서드에서는 다음과 같이 처리합니다.

• 루트(root) 노드가 비어 있으면(빈 트리) 0을 반환합니다.
• 그렇지 않으면 recurse(root)의 결과를 반환합니다.

구현 예제

아래 코드를 통해 더 자세히 이해해 보겠습니다.

class TreeNode:
   def __init__(self, data, left = None, right = None):
      self.val = data
      self.left = left
      self.right = right
class Solution:
   def recurse(self, node):
      val = node.val
      if node.left:
         val += self.recurse(node.left)
      if node.right:
         val += self.recurse(node.right)
      return val
   def solve(self, root):
      if not root:
         return 0
      return self.recurse(root)
ob = Solution()
root = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(3)
root.right.right = TreeNode(5)
print(ob.solve(root))

입력

root = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(3)
root.right.right = TreeNode(5)

출력

14

복잡도 분석

이 알고리즘은 트리의 모든 노드를 한 번씩 방문하므로 시간 복잡도는 O(n)입니다. 여기서 n은 트리의 노드 개수입니다. 공간 복잡도는 재귀 호출 스택의 깊이에 의해 결정되며, 균형 잡힌 트리의 경우 O(log n), 최악의 경우(편향된 트리) O(n)입니다.