Python으로 이진 트리에서 자식이 하나뿐인 노드 모두 제거하는 방법

이진 트리의 루트(root)가 주어졌을 때, 자식이 하나뿐인 노드를 모두 제거하는 문제를 살펴보겠습니다. 즉, 왼쪽 또는 오른쪽 자식 중 하나만 가지고 있는 노드는 트리에서 삭제하고, 리프 노드와 두 자식을 모두 가진 노드는 그대로 유지해야 합니다.

문제 예시

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

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

풀이 접근 방식

이 문제는 재귀(Recursion)를 활용하면 간단하게 해결할 수 있습니다. 각 노드를 방문하면서 자식 수를 확인하고, 조건에 맞지 않는 노드는 하위 트리로 치환하는 방식입니다.

구체적인 알고리즘은 다음과 같습니다.

• solve()라는 메서드를 정의하고, 트리의 루트를 인자로 전달받습니다.
• 루트가 null이면 그대로 루트를 반환합니다.
• 루트의 왼쪽과 오른쪽 자식이 모두 null이라면(리프 노드), 해당 노드를 반환합니다.
• 루트의 왼쪽 자식이 null이라면, 오른쪽 서브트리에 대해 solve()를 재귀 호출한 결과를 반환합니다.
• 루트의 오른쪽 자식이 null이라면, 왼쪽 서브트리에 대해 solve()를 재귀 호출한 결과를 반환합니다.
• 그 외의 경우에는 왼쪽과 오른쪽 서브트리 각각에 대해 solve()를 재귀 호출한 결과를 연결한 뒤 루트를 반환합니다.

예제 코드

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

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

class Solution:
   def solve(self, root):
      # 빈 노드인 경우 그대로 반환
      if not root:
         return root

      # 리프 노드인 경우 그대로 유지
      if not root.left and not root.right:
         return root

      # 왼쪽 자식이 없으면 오른쪽 서브트리로 대체
      if not root.left:
         return self.solve(root.right)

      # 오른쪽 자식이 없으면 왼쪽 서브트리로 대체
      if not root.right:
         return self.solve(root.left)

      # 양쪽 자식이 모두 있으면 각 서브트리를 재귀적으로 정리
      root.left = self.solve(root.left)
      root.right = self.solve(root.right)

      return root

ob = Solution()
root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.right.right = TreeNode(5)
root.left.left.right = TreeNode(6)
root.right.right.left = TreeNode(7)
root.right.right.right = TreeNode(8)

res = ob.solve(root)
print_tree(res)

입력

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.right.right = TreeNode(5)
root.left.left.right = TreeNode(6)
root.right.right.left = TreeNode(7)
root.right.right.right = TreeNode(8)

출력

6, 1, 7, 5, 8,

동작 원리 정리

위 예제에서 노드 2와 4는 자식이 하나뿐이므로 제거되고, 해당 위치에는 실질적인 리프 노드인 6이 올라오게 됩니다. 마찬가지로 노드 3도 자식이 하나뿐이므로 제거되어 노드 5가 루트의 오른쪽 자리를 차지합니다. 최종 결과는 중위 순회(inorder traversal) 기준으로 6, 1, 7, 5, 8이 출력됩니다.

이 알고리즘의 시간 복잡도는 트리의 모든 노드를 한 번씩 방문하므로 O(n)이며, 재귀 호출 깊이에 따라 공간 복잡도는 트리의 높이인 O(h)입니다.