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

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

지그재그 레벨 순회(Zigzag Level Order Traversal)란?

이진 트리가 주어졌을 때, 지그재그 레벨 순회는 트리를 레벨(층) 단위로 방문하되 방향을 매 층마다 번갈아 바꾸는 순회 기법입니다. 첫 번째 레벨은 왼쪽에서 오른쪽으로, 두 번째 레벨은 오른쪽에서 왼쪽으로, 세 번째 레벨은 다시 왼쪽에서 오른쪽으로 탐색하는 방식으로 진행됩니다.

예를 들어 아래와 같은 이진 트리가 있다고 가정해 보겠습니다.

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

이 트리를 지그재그 레벨 순회하면 결과는 [[3], [20, 9], [15, 7]]이 됩니다.

  • 레벨 0: [3] → 왼쪽에서 오른쪽으로
  • 레벨 1: [20, 9] → 오른쪽에서 왼쪽으로
  • 레벨 2: [15, 7] → 왼쪽에서 오른쪽으로

문제 해결 접근 방법

이 문제는 큐(queue)와 현재 순회 방향을 나타내는 불리언 플래그(flag)를 활용하면 깔끔하게 해결할 수 있습니다. 전체 알고리즘은 다음과 같습니다.

  1. 트리가 비어 있으면 빈 리스트를 반환합니다.
  2. 큐를 생성하고 루트 노드를 삽입합니다. 노드 객체를 저장할 리스트 res, 최종 결과값을 저장할 리스트 res2를 준비하고, 방향 플래그를 True(왼쪽→오른쪽)로 설정합니다.
  3. 큐가 빌 때까지 다음 과정을 반복합니다.
    • 현재 큐에 들어 있는 노드들을 res에 추가하고, 해당 노드들의 값을 res2에 추가합니다.
    • flag가 True라면, 현재 레벨의 노드를 뒤에서부터 앞으로 순회하면서 오른쪽 자식을 먼저, 그다음 왼쪽 자식을 새 큐에 삽입합니다. 이렇게 하면 다음 레벨이 자연스럽게 오른쪽→왼쪽 순서로 채워집니다.
    • flag가 False라면 반대로 왼쪽 자식을 먼저 삽입하여 다음 레벨이 왼쪽→오른쪽 순서가 되도록 합니다.
    • 플래그 값을 반전시켜 다음 레벨의 탐색 방향을 전환합니다.
  4. 모든 레벨의 순회가 끝나면 res2를 반환합니다.

파이썬 구현 예제

아래 코드는 위에서 설명한 알고리즘을 파이썬으로 구현한 전체 예제입니다.

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


def insert(temp, data):
    # 트리에 값을 레벨 순서(BFS)로 삽입하는 헬퍼 함수
    if data is None:
        return
    que = [temp]
    while que:
        temp = que.pop(0)
        if not temp.left:
            temp.left = TreeNode(data)
            break
        else:
            que.append(temp.left)
        if not temp.right:
            temp.right = TreeNode(data)
            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 zigzagLevelOrder(self, root):
        # 트리가 비어 있는 경우 빈 리스트 반환
        if not root:
            return []

        queue = [root]
        res = []       # 각 레벨의 노드 객체를 저장
        res2 = []      # 최종 결과(각 레벨의 값 리스트)
        flag = True    # True: 왼쪽→오른쪽, False: 오른쪽→왼쪽

        while queue:
            res.append(list(queue))
            res2.append([node.data for node in queue])

            next_queue = []
            if flag:
                # 현재 레벨을 뒤에서부터 순회하며 오른쪽 자식을 먼저 삽입
                for i in range(len(res[-1]) - 1, -1, -1):
                    if res[-1][i].right:
                        next_queue.append(res[-1][i].right)
                    if res[-1][i].left:
                        next_queue.append(res[-1][i].left)
            else:
                # 현재 레벨을 뒤에서부터 순회하며 왼쪽 자식을 먼저 삽입
                for i in range(len(res[-1]) - 1, -1, -1):
                    if res[-1][i].left:
                        next_queue.append(res[-1][i].left)
                    if res[-1][i].right:
                        next_queue.append(res[-1][i].right)

            queue = next_queue
            flag = not flag

        return res2


ob = Solution()
tree = make_tree([3, 9, 20, None, None, 15, 7])
print(ob.zigzagLevelOrder(tree))

실행 결과 확인

입력

[3, 9, 20, null, null, 15, 7]

출력

[[3], [20, 9], [15, 7]]

시간 및 공간 복잡도

  • 시간 복잡도: O(n) — 트리의 모든 노드를 정확히 한 번씩 방문합니다.
  • 공간 복잡도: O(n) — 큐와 결과 리스트에 트리의 노드 수만큼 공간이 필요합니다.

지그재그 레벨 순회는 일반적인 레벨 순회(BFS) 코드에 방향 전환 로직만 추가하면 되기 때문에, 레벨 순회의 원리만 잘 이해하고 있다면 어렵지 않게 구현할 수 있습니다. 코딩 테스트에서 자주 등장하는 유형이므로 큐와 플래그를 조합하는 이 패턴을 꼭 익혀두시기 바랍니다.