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

파이썬으로 구현하는 이진 트리 지그재그 레벨 순회 (Zigzag Level Order Traversal)

문제 소개

이진 트리가 하나 주어졌다고 가정해 보겠습니다. 이때 각 레벨의 노드 값들을 첫 번째 레벨은 왼쪽에서 오른쪽, 다음 레벨은 오른쪽에서 왼쪽 방향으로 번갈아 가며 순회한 결과를 출력해야 합니다. 이러한 순회 방식은 지그재그(Zigzag) 레벨 순회 또는 나선형(Spiral) 순회라고도 불립니다.

예를 들어 입력 트리가 다음과 같다면,

파이썬으로 구현하는 이진 트리 지그재그 레벨 순회 (Zigzag Level Order Traversal)

출력 결과는 다음과 같습니다.

[5, -10, 4, -2, -7, 15]

해결 전략: 두 개의 스택 활용

이 문제는 두 개의 스택(stack)을 사용하면 효율적으로 해결할 수 있습니다. 스택의 LIFO(Last-In-First-Out) 특성 덕분에 레벨마다 탐색 방향을 자연스럽게 뒤집을 수 있기 때문입니다. 알고리즘의 단계별 동작은 다음과 같습니다.

  • 루트가 null이면 빈 리스트를 반환합니다.
  • s1 := 루트 노드를 초기값으로 담은 리스트(첫 번째 스택)
  • s2 := 비어 있는 새로운 리스트(두 번째 스택)
  • res := 결과 값을 저장할 새로운 리스트
  • s1 또는 s2가 비어 있지 않은 동안 다음을 반복합니다.
    • s1 처리 (왼쪽 → 오른쪽 레벨):
      • node := s1에서 마지막 요소를 꺼냅니다(pop).
      • 노드의 왼쪽 자식이 존재하면 s2의 끝에 추가합니다.
      • 노드의 오른쪽 자식이 존재하면 s2의 끝에 추가합니다.
      • 노드의 값을 res의 끝에 추가합니다.
    • s2 처리 (오른쪽 → 왼쪽 레벨):
      • node := s2에서 마지막 요소를 꺼냅니다(pop).
      • 노드의 오른쪽 자식이 존재하면 s1의 끝에 추가합니다.
      • 노드의 왼쪽 자식이 존재하면 s1의 끝에 추가합니다.
      • 노드의 값을 res의 끝에 추가합니다.
  • 모든 노드를 순회한 후 res를 반환합니다.

여기서 핵심 포인트는 두 스택에 자식 노드를 삽입하는 순서가 서로 반대라는 점입니다. s1에서는 왼쪽 자식을 먼저 넣지만, s2에서는 오른쪽 자식을 먼저 넣습니다. 이렇게 하면 다음 레벨에서 pop할 때 원하는 방향 순서대로 노드가 꺼내지게 됩니다.

구현 예제

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

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

class Solution:
    def solve(self, root):
        if not root:
            return []
        s1 = [root]
        s2 = []
        res = []
        while s1 or s2:
            while s1:
                node = s1.pop()
                if node.left:
                    s2.append(node.left)
                if node.right:
                    s2.append(node.right)
                res.append(node.val)
            while s2:
                node = s2.pop()
                if node.right:
                    s1.append(node.right)
                if node.left:
                    s1.append(node.left)
                res.append(node.val)
        return res

ob = Solution()
root = TreeNode(5)
root.left = TreeNode(4)
root.right = TreeNode(-10)
root.left.left = TreeNode(-2)
root.right.left = TreeNode(-7)
root.right.right = TreeNode(15)
print(ob.solve(root))

입력

root = TreeNode(5)
root.left = TreeNode(4)
root.right = TreeNode(-10)
root.left.left = TreeNode(-2)
root.right.left = TreeNode(-7)
root.right.right = TreeNode(15)

출력

[5, -10, 4, -2, -7, 15]

복잡도 분석

시간 복잡도: O(n) — 모든 노드를 정확히 한 번씩 방문합니다.
공간 복잡도: O(n) — 최악의 경우 한 레벨의 모든 노드가 스택에 저장될 수 있습니다.