파이썬으로 이진 트리에서 두 노드 사이의 거리 구하기

이진 트리가 주어졌을 때, 그 안에서 두 노드 사이의 거리를 찾아야 하는 상황을 가정해 보겠습니다. 그래프에서처럼 두 노드를 연결하는 간선(edge)들을 따라가며, 그 경로에 있는 간선의 개수를 세어 거리로 반환하면 됩니다. 이때 트리의 노드는 다음과 같은 구조를 가집니다.

data : <정수 값>
right : <트리의 다른 노드를 가리키는 포인터>
left : <트리의 다른 노드를 가리키는 포인터>

예를 들어 입력이 다음과 같다고 해보겠습니다.

거리를 구해야 할 두 노드가 2와 8이라면, 출력 결과는 4가 됩니다.

노드 2와 8 사이의 간선은 (2, 3), (3, 5), (5, 7), (7, 8)로 총 4개입니다. 따라서 두 노드 사이의 거리는 4입니다.

문제 해결 접근 방법

이 문제는 최소 공통 조상(LCA, Lowest Common Ancestor)을 활용하면 효율적으로 해결할 수 있습니다. 두 노드의 거리는 '각 노드에서 LCA까지의 거리'를 더한 값과 같기 때문입니다. 구체적인 단계는 다음과 같습니다.

findLca() 함수를 정의합니다. 이 함수는 root, p, q를 인자로 받습니다.
- root가 null이면 null을 반환합니다.
- root의 데이터가 p 또는 q 중 하나와 같다면 root를 반환합니다.
- left := findLca(root의 왼쪽 자식, p, q)
- right := findLca(root의 오른쪽 자식, p, q)
- left와 right가 모두 null이 아니라면 root가 곧 LCA이므로 root를 반환합니다.
- 그렇지 않으면 left 또는 right를 반환합니다.
findDist() 함수를 정의합니다. 이 함수는 root, data를 인자로 받으며, BFS(너비 우선 탐색)를 사용해 특정 노드까지의 거리를 구합니다.
- queue := 새로운 deque를 생성합니다.
- queue의 끝에 (root, 0) 쌍을 삽입합니다.
- queue가 비어 있지 않은 동안 다음을 반복합니다.
  - current := queue 맨 앞 쌍의 첫 번째 값
  - dist := queue 맨 앞 쌍의 두 번째 값
  - current의 데이터가 data와 같다면 dist를 반환합니다.
  - current의 왼쪽 자식이 null이 아니라면 (왼쪽 자식, dist+1)을 queue에 추가합니다.
  - current의 오른쪽 자식이 null이 아니라면 (오른쪽 자식, dist+1)을 queue에 추가합니다.
node := findLca(root, p, q)로 LCA를 구합니다.
findDist(node, p) + findDist(node, q)를 반환합니다. 즉, LCA에서 각 노드까지의 거리를 합산한 값이 곧 두 노드 사이의 거리입니다.

예제 코드

다음 파이썬 구현을 통해 동작 방식을 더 잘 이해해 보겠습니다.

import collections
class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.data = data
        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.data == element):
        return root

    res1 = search_node(root.left, element)
    if res1:
        return res1

    res2 = search_node(root.right, element)
    return res2

def print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.data, end == ', ')
        print_tree(root.right)

def findLca(root, p, q):
    if root is None:
        return None
    if root.data in (p,q):
        return root
    left = findLca(root.left, p, q)
    right = findLca(root.right, p, q)
    if left and right:
        return root
    return left or right

def findDist(root, data):
    queue = collections.deque()
    queue.append((root, 0))
    while queue:
        current, dist = queue.popleft()
        if current.data == data:
            return dist
        if current.left: queue.append((current.left, dist+1))
        if current.right: queue.append((current.right, dist+1))

def solve(root, p, q):
    node = findLca(root, p, q)
    return findDist(node, p) + findDist(node, q)

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

입력

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