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

Python으로 풀이하는 이진 트리 경로 합(Path Sum)

문제 개요

하나의 이진 트리와 목표 합(sum)이 주어졌다고 가정해 보겠습니다. 우리가 찾아야 하는 것은 루트 노드에서 시작하여 리프 노드까지 따라 내려가는 경로 중, 경로상 노드 값들의 합이 주어진 값과 정확히 일치하는 경로입니다.

예를 들어 트리가 [0, -3, 9, -10, null, 5]이고 목표 합이 14라고 한다면, 0 → 9 → 5 경로를 따라가면 0 + 9 + 5 = 14가 되므로 조건을 만족하는 경로가 존재합니다.

Python으로 풀이하는 이진 트리 경로 합(Path Sum)

해결 접근 방법

이 문제는 재귀(DFS) 방식으로 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 루트에서 리프로 내려갈 때마다 현재 노드의 값을 목표 합에서 차감하고, 리프 노드에 도달했을 때 남은 값이 0인지 확인하는 것입니다. 단계별로 살펴보면 다음과 같습니다.

  • 루트가 null이면 탐색할 경로가 없으므로 False를 반환합니다.
  • 현재 노드가 리프 노드(왼쪽과 오른쪽 자식이 모두 없는 노드)라면, sum - root.val == 0일 때 true를 반환하고, 그렇지 않으면 false를 반환합니다.
  • 그 외의 경우에는 왼쪽 서브트리와 오른쪽 서브트리 각각에 대해 남은 합(sum - root.val)을 전달하며 재귀적으로 탐색한 뒤, 두 결과를 OR 연산으로 결합하여 반환합니다. 어느 한쪽이라도 유효한 경로를 찾으면 true가 됩니다.

시간 복잡도는 모든 노드를 최악의 경우 한 번씩 방문하므로 O(n)이며, 공간 복잡도는 재귀 호출 스택의 깊이에 비례하여 균형 잡힌 트리 기준 O(log n), 편향된 트리의 경우 O(n)입니다.

구현 예제

아래 구현을 통해 더 자세히 이해해 보겠습니다.

# 이진 트리 노드 정의
class TreeNode(object):
    def __init__(self, x):
        self.data = x
        self.left = None
        self.right = None

def insert(temp, data):
    que = []
    que.append(temp)
    while (len(que)):
        temp = que[0]
        que.pop(0)
        if (not temp.left):
            if data is not None:
                temp.left = TreeNode(data)
            else:
                temp.left = TreeNode(0)
            break
        else:
            que.append(temp.left)
        if (not temp.right):
            if data is not None:
                temp.right = TreeNode(data)
            else:
                temp.right = TreeNode(0)
            break
        else:
            que.append(temp.right)

def make_tree(elements):
    Tree = TreeNode(elements[0])
    for element in elements[1:]:
        insert(Tree, element)
    return Tree

class Solution(object):
    def hasPathSum(self, root, sum):
        """
        :type root: TreeNode
        :type sum: int
        :rtype: bool
        """
        if not root:
            return False
        if not root.left and not root.right and root.data is not None:
            return sum - root.data == 0
        if root.data is not None:
            return self.hasPathSum(root.left, sum-root.data) or self.hasPathSum(root.right, sum-root.data)

tree1 = make_tree([0,-3,9,-10,None,5])
ob1 = Solution()
print(ob1.hasPathSum(tree1, 14))

입력

tree1 = make_tree([0,-3,9,-10,None,5])

출력

True

동작 원리 정리

위 코드에서 hasPathSum 메서드는 먼저 빈 트리 여부를 검사합니다. 이후 현재 노드가 리프 노드인지 확인하여, 리프라면 지금까지 차감해 온 목표 합에서 마지막 노드 값을 뺀 결과가 0인지 판단합니다. 리프가 아니라면 왼쪽과 오른쪽 자식에게 각각 줄어든 목표 합을 넘겨주며 재귀 호출을 수행하고, 두 탐색 결과 중 하나라도 true이면 해당 경로가 존재한다는 의미이므로 true를 반환합니다. 이처럼 목표 합을 위에서 아래로 점진적으로 차감해 나가는 방식 덕분에 별도의 누적 변수 없이도 간결하게 문제를 해결할 수 있습니다.