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

파이썬으로 이진 트리의 최소 공통 조상(LCA) 찾기: DFS 재귀 탐색 구현

이진 트리와 두 개의 특정 노드 x, y가 주어졌다고 가정해 봅시다. 우리가 찾아야 하는 것은 이진 트리 안에서 이 두 노드의 최소 공통 조상(Lowest Common Ancestor, LCA)입니다.

이진 트리에서 최소 공통 조상이란, 노드 x와 y가 모두 그 노드의 자손(descendant)에 해당하는 노드 중에서 가장 깊은 위치에 있는 노드를 의미합니다. 참고로, 하나의 노드는 자기 자신의 자손이 될 수도 있습니다. 즉, x가 y의 조상이라면 x 자체가 최소 공통 조상이 될 수 있습니다.

예를 들어 아래와 같은 트리가 있고 x = 2, y = 4라고 한다면, 출력 결과는 3이 됩니다.

파이썬으로 이진 트리의 최소 공통 조상(LCA) 찾기: DFS 재귀 탐색 구현

노드 2와 노드 4가 모두 자손인 가장 낮은 노드는 3이기 때문입니다. 따라서 3을 반환하면 됩니다.

해결 접근 방법: DFS 재귀 탐색

이 문제는 깊이 우선 탐색(Depth-First Search, DFS)을 이용하면 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 트리를 순회하면서 두 목표 노드가 어느 지점에서 서로 다른 하위 트리로 갈라지는지, 혹은 어느 노드가 두 노드를 모두 포함하는지 확인하는 것입니다.

단계별로 살펴보면 다음과 같습니다.

  1. dfs() 함수 정의 — 노드를 인자로 받는 재귀 함수를 정의합니다.
  2. null 체크 — 노드가 null이면 아무것도 반환하지 않고 종료합니다.
  3. 목표 노드 발견 시 — 현재 노드가 [x, y] 목록에 있다면, 왼쪽과 오른쪽 자식에 대해 dfs()를 재귀 호출합니다. 그중 하나라도 null이 아니면(즉, 다른 목표 노드가 하위 트리에 존재하면) 현재 노드가 최소 공통 조상이므로 ans에 저장하고 반환합니다.
  4. 일반 노드인 경우 — 왼쪽과 오른쪽 자식에 대해 각각 dfs()를 재귀 호출합니다.
  5. 분기점 판별 — 양쪽 탐색 결과가 모두 null이 아니라면, 두 목표 노드가 서로 다른 하위 트리에 존재한다는 의미이므로 현재 노드가 최소 공통 조상입니다. ans에 저장하고 반환합니다.
  6. 결과 전달 — 위 경우가 아니라면 null이 아닌 쪽의 결과(left 또는 right)를 상위 호출로 반환합니다.
  7. 최종 호출 — 루트 노드부터 dfs()를 호출하고 그 결과를 ans에 담아 반환합니다.

파이썬 구현 예제

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 is 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 solve(root, x, y):
   def dfs(node):
      if not node:
         return
      if node in [x, y]:
         left = dfs(node.left)
         right = dfs(node.right)
         if left or right:
            return node
      left = dfs(node.left)
      right = dfs(node.right)
      if left and right:
         return node
      return left or right

   ans = dfs(root)
   return ans


root = make_tree([5, 3, 7, 2, 4, 1, 7, 6, 8, 10])
print(solve(root, search_node(root, 2), search_node(root, 4)).data)

입력

make_tree([5, 3, 7, 2, 4, 1, 7, 6, 8, 10]),
search_node(root, 2),
search_node(root, 4)

출력

3

코드 설명

위 코드는 크게 세 부분으로 구성됩니다.

  • TreeNode 클래스 — 트리의 각 노드를 나타내며, 데이터 값과 왼쪽·오른쪽 자식 노드를 저장합니다.
  • insert(), make_tree(), search_node() — 리스트 형태의 입력값으로 이진 트리를 생성하고(make_tree), 특정 값을 가진 노드를 찾아내는(search_node) 보조 함수들입니다.
  • solve() 함수 — 내부에 정의된 dfs() 재귀 함수를 통해 최소 공통 조상을 실제로 찾아내는 핵심 로직입니다.

dfs() 함수는 트리를 후위 순회(postorder) 방식으로 탐색합니다. 자식 노드들의 탐색 결과를 먼저 확인한 뒤 현재 노드를 판단하기 때문에, 두 목표 노드가 처음으로 만나는 지점, 즉 최소 공통 조상을 정확하게 찾아낼 수 있습니다.

시간 복잡도는 트리의 모든 노드를 한 번씩 방문하므로 O(n)이며, 공간 복잡도는 재귀 호출 스택의 깊이에 비례하여 최악의 경우(편향된 트리) O(n)입니다.