Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 이진 트리에서 루트부터 리프까지 가장 긴 경로의 합 구하기

문제 소개

이진 트리가 하나 주어져 있을 때, 루트(root) 노드에서 리프(leaf) 노드까지 이어지는 가장 긴 경로의 노드 값 합계를 구하는 것이 목표입니다. 만약 길이가 같은 경로가 둘 이상 존재한다면, 그중 합이 더 큰 경로의 값을 반환해야 합니다.

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

파이썬으로 이진 트리에서 루트부터 리프까지 가장 긴 경로의 합 구하기

출력 결과는 20이 됩니다. 루트 2에서 시작해 4 → 8 → 6으로 이어지는 경로가 가장 길며, 그 합은 2 + 4 + 8 + 6 = 20이기 때문입니다.

풀이 접근 방식

이 문제는 재귀(recursion)를 활용하면 간단하게 해결할 수 있습니다. 각 노드마다 "현재 노드를 루트로 하는 서브트리에서 가장 긴 경로의 길이"와 "그 경로의 노드 값 합"을 함께 반환하는 방식입니다.

  1. rec() 함수 정의: 이 함수는 현재 노드 curr를 인자로 받습니다.
  2. 기저 조건: curr가 null(None)이면 (0, 0)을 반환합니다. 즉, 경로 길이와 합이 모두 0입니다.
  3. 자식 노드 비교: bigger := rec(curr의 왼쪽 자식)과 rec(curr의 오른쪽 자식) 반환값 중 더 큰 값으로 설정합니다.
  4. 결과 반환: (bigger[0] + 1, bigger[1] + curr.val) 쌍을 반환합니다. 첫 번째 원소는 경로 길이, 두 번째 원소는 경로의 합입니다.
  5. 메인 로직: 메인 메서드에서 ret := rec(root)를 호출한 뒤, ret의 두 번째 원소(인덱스 1), 즉 가장 긴 경로의 합을 반환합니다.

여기서 주목할 점은 파이썬의 max() 함수가 튜플을 사전순(lexicographic order)으로 비교한다는 것입니다. 따라서 먼저 경로 길이를 비교하고, 길이가 같을 경우에는 합을 기준으로 더 큰 쪽을 자동으로 선택합니다. 덕분에 "길이가 같으면 합이 큰 경로를 반환한다"는 조건을 별도의 처리 없이 자연스럽게 만족할 수 있습니다.

예제 코드

class TreeNode:
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def solve(self, root):
        def rec(curr):
            if not curr:
                return (0, 0)
            bigger = max(rec(curr.left), rec(curr.right))
            return (bigger[0] + 1, bigger[1] + curr.val)
        return rec(root)[1]

ob = Solution()
root = TreeNode(2)
root.left = TreeNode(10)
root.right = TreeNode(4)
root.right.left = TreeNode(8)
root.right.right = TreeNode(2)
root.right.left.left = TreeNode(6)
print(ob.solve(root))

입력

root = TreeNode(2)
root.left = TreeNode(10)
root.right = TreeNode(4)
root.right.left = TreeNode(8)
root.right.right = TreeNode(2)
root.right.left.left = TreeNode(6)

출력

20

정리

이 풀이는 모든 노드를 한 번씩만 방문하므로 시간 복잡도는 O(N)(N은 노드 수)이며, 재귀 호출 스택의 깊이만큼 추가 공간이 필요하므로 공간 복잡도는 O(H)(H는 트리의 높이)입니다. 튜플 비교를 활용한 간결한 재귀 구조 덕분에 "가장 긴 경로를 찾되, 같은 길이라면 더 큰 합을 반환한다"는 조건을 우아하게 처리할 수 있습니다.