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

파이썬으로 이진 트리에서 가장 긴 지그재그 경로의 길이 찾기


문제 소개

이진 트리가 하나 주어졌을 때, 왼쪽 자식과 오른쪽 자식을 번갈아 이동하면서 아래 방향으로만 내려가는 경로 중 가장 긴 것의 길이를 찾아야 합니다. 이러한 경로는 일반적으로 지그재그 경로(ZigZag Path)라고 불립니다.

예를 들어 아래와 같은 트리가 입력으로 주어진 경우를 살펴보겠습니다.

파이썬으로 이진 트리에서 가장 긴 지그재그 경로의 길이 찾기

이때 정답은 5입니다. [2 → 4 → 5 → 7 → 8] 경로가 오른쪽 → 왼쪽 → 오른쪽 → 왼쪽으로 방향을 번갈아 가며 내려가는 가장 긴 교대 경로이기 때문입니다.

풀이 접근 방법

이 문제는 깊이 우선 탐색(DFS)으로 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 각 노드를 방문할 때 '다음 이동이 왼쪽인지 오른쪽인지'를 나타내는 플래그(flag)를 함께 전달하여, 경로가 계속되는 방향의 자식에서는 길이를 1씩 늘려 이어가고, 반대편 자식에서는 길이 1부터 새로운 경로를 시작하는 것입니다.

구체적인 단계는 다음과 같습니다.

  • 루트가 null이면 0을 반환합니다.
  • dfs(node, count, flag) 함수를 정의합니다. flag는 다음에 왼쪽으로 이동할 차례인지(True), 오른쪽으로 이동할 차례인지(False)를 나타냅니다.
  • node가 null이 아니라면 다음을 수행합니다.
    • flag가 True라면(왼쪽으로 경로를 이어감):
      • a := dfs(node의 왼쪽 자식, count + 1, False)
      • b := dfs(node의 오른쪽 자식, 1, True)
    • flag가 False라면(오른쪽으로 경로를 이어감):
      • a := dfs(node의 오른쪽 자식, count + 1, True)
      • b := dfs(node의 왼쪽 자식, 1, False)
    • a와 b 중 최댓값을 반환합니다.
  • node가 null이면 지금까지 쌓인 count를 반환합니다.
  • 메인 로직에서는 다음을 수행합니다.
    • a := dfs(루트의 왼쪽 자식, 1, False)
    • b := dfs(루트의 오른쪽 자식, 1, True)
    • a와 b 중 최댓값을 최종 결과로 반환합니다.

여기서 count + 1로 호출하는 경우는 현재 진행 중인 교대 경로를 이어가는 것이고, count를 1로 초기화하여 호출하는 경우는 해당 자식 노드에서 새로운 교대 경로를 시작하는 것을 의미합니다.

파이썬 구현 예제

다음 구현을 통해 풀이 과정을 더 잘 이해할 수 있습니다.

예제 코드

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

class Solution:
    def solve(self, root):
        if not root:
            return 0

        def dfs(node, count, flag):
            if node:
                if flag == True:
                    a = dfs(node.left, count + 1, False)
                    b = dfs(node.right, 1, True)
                elif flag == False:
                    a = dfs(node.right, count + 1, True)
                    b = dfs(node.left, 1, False)
                return max(a, b)
            return count

        a = dfs(root.left, 1, False)
        b = dfs(root.right, 1, True)
        return max(a, b)

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

입력

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

출력

5

복잡도 분석

  • 시간 복잡도: O(N) — 각 노드가 상수 시간의 작업과 함께 한 번씩 처리됩니다.
  • 공간 복잡도: O(H) — 재귀 호출 스택이 트리의 높이(H)만큼 사용되며, 기울어진(skewed) 트리의 경우 최악에 O(N)이 됩니다.