Python으로 이진 트리에서 특정 노드의 오른쪽 노드 찾는 방법

이진 트리(binary tree)가 하나 주어지고, 특정 노드를 가리키는 포인터 u도 함께 제공된다고 가정해 봅시다. 우리가 해야 할 일은 이 노드 바로 오른쪽에 있는 노드를 찾는 것입니다.

여기서 중요한 조건은 다음과 같습니다.

• 오른쪽에 있는 노드는 반드시 같은 레벨(level), 즉 같은 깊이에 위치해야 합니다.
• 주어진 노드 u는 리프 노드일 수도 있고 내부 노드일 수도 있습니다.

문제 예시

예를 들어 다음과 같은 트리가 있다고 가정합니다.

root = make_tree([5, 3, 7, 2, 4, 6, 8])

이때 u = 6이라면 출력 결과는 8이 됩니다. 노드 6과 같은 레벨에서 바로 오른쪽에 있는 노드가 8이기 때문입니다.

접근 방법: 레벨 순회(BFS) 활용

이 문제는 BFS(너비 우선 탐색) 기반의 레벨 순회로 해결할 수 있습니다. 핵심 아이디어는 트리를 레벨별로 순회하면서, 목표 노드 u가 속한 레벨에서 그다음 위치에 있는 노드를 찾는 것입니다.

알고리즘 단계

1. 루트(root)가 비어 있다면 null(None)을 반환합니다.
2. 새로운 deque(덱) dq를 생성하고, 루트를 넣습니다.
3. dq가 빌 때까지 다음을 반복합니다.
   • dq_size := 현재 dq의 크기 (현재 레벨의 노드 수)
   • temp := 현재 레벨의 노드들을 저장할 새 리스트
   • index := -1 (목표 노드의 위치)
   • 0부터 dq_size까지 반복하면서:
     • node := dq에서 요소를 꺼냄
     • node의 왼쪽 자식이 존재하면 dq에 추가
     • node의 오른쪽 자식이 존재하면 dq에 추가
     • temp 리스트의 끝에 node를 추가
     • 만약 node가 u와 같다면, index := len(temp) - 1
   • 반복 종료 후, index가 temp의 마지막 인덱스와 같다면 u가 해당 레벨의 가장 오른쪽 노드이므로 null을 반환합니다.
   • index가 -1보다 크다면, temp[index + 1], 즉 바로 오른쪽 노드를 반환합니다.
4. 모든 레벨을 확인했는데도 찾지 못했다면 null을 반환합니다.

Python 구현 코드

아래는 위 알고리즘을 실제로 구현한 전체 코드입니다.

from queue import deque
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

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

def search_node(root, element):
    if (root == None):
        return None
    if (root.val == element):
        return root
    res1 = search_node(root.left, element)
    if res1:
        return res1
    res2 = search_node(root.right, element)
    return res2

def solve(root, u):
    if not root:
        return None
    dq = deque()
    dq.append(root)
    while dq:
        dq_size = len(dq)
        temp = []
        index = -1
        for _ in range(dq_size):
            node = dq.pop()
            if node.left:
                dq.appendleft(node.left)
            if node.right:
                dq.appendleft(node.right)
            temp.append(node)
            if node == u:
                index = len(temp) - 1
        if index == len(temp) - 1:
            return None
        if index > -1:
            return temp[index + 1]
    return None

root = make_tree([5, 3, 7, 2, 4, 6, 8])
u = search_node(root, 6)
ret = solve(root, u)
print(ret.val)

입력

root = make_tree([5, 3, 7, 2, 4, 6, 8])
u = search_node(root, 6)

출력

8

정리

이 알고리즘은 트리의 모든 노드를 한 번씩 방문하므로 시간 복잡도는 O(n)이며, 최악의 경우 한 레벨의 모든 노드를 저장해야 하므로 공간 복잡도 역시 O(n)입니다. 만약 목표 노드가 해당 레벨의 마지막 노드라면 오른쪽에 있는 노드가 존재하지 않으므로 None이 반환됩니다.