파이썬으로 트리를 두 개로 분할하는 경우의 수 세기

값이 0, 1, 2로만 구성된 이진 트리(binary tree)가 있다고 가정해 봅시다. 루트 노드에는 최소한 하나의 0 노드와 하나의 1 노드가 존재합니다. 이제 트리에서 간선(edge) 하나를 삭제하면 트리가 서로 다른 두 개의 트리로 분리되는 연산이 있다고 합시다. 우리가 구해야 할 것은, 분리된 두 트리 어느 쪽에도 0 노드와 1 노드가 동시에 포함되지 않도록 간선을 삭제할 수 있는 경우의 수입니다.

예를 들어 입력이 아래와 같은 트리라면,

출력은 1이 됩니다. 0 노드와 2 노드 사이의 간선만이 조건을 만족하며 삭제 가능하기 때문입니다.

문제 해결 접근 방법

이 문제는 다음 단계를 따라 해결할 수 있습니다.

• 전역 카운터 count := [0, 0, 0]를 초기화합니다.
• dfs() 함수를 정의합니다. 이 함수는 노드를 인자로 받습니다.
  • 노드가 null이 아니라면:
    • pre := count로 현재 카운트를 백업합니다.
    • 왼쪽 자식과 오른쪽 자식에 대해 재귀적으로 dfs()를 호출합니다.
    • count[노드의 값] += 1로 현재 노드의 값을 반영합니다.
    • node.count에 해당 노드의 서브트리에 포함된 0과 1의 개수, 즉 (count[i] - pre[i]) (i = 0, 1)를 저장합니다.
• dfs2() 함수를 정의합니다. 이 함수는 노드와 부모 노드(par)를 인자로 받습니다.
  • 노드가 null이 아니라면:
    • 부모 노드가 null이 아니라면:
      • (a0, a1) := 노드의 count — 노드 쪽 서브트리의 0, 1 개수입니다.
      • (b0, b1) := (count[0] - a0, count[1] - a1) — 나머지 부분(부모 쪽 트리)의 0, 1 개수입니다.
      • 만약 (a0 == 0 또는 a1 == 0) 그리고 (b0 == 0 또는 b1 == 0)이라면, 즉 양쪽 모두 0과 1이 동시에 존재하지 않는다면:
        • ans += 1로 정답 카운트를 증가시킵니다.
    • 왼쪽 자식과 오른쪽 자식에 대해 dfs2()를 재귀 호출합니다.
• 메인 흐름에서는 다음을 수행합니다.
  • dfs(root)를 호출해 각 노드의 서브트리 통계를 계산합니다.
  • ans := 0으로 초기화합니다.
  • dfs2(root)를 호출해 유효한 분할 지점을 셉니다.
  • ans를 반환합니다.

구현 예제

아래 코드를 통해 더 잘 이해할 수 있습니다.

class TreeNode:
   def __init__(self, data, left = None, right = None):
      self.val = data
      self.left = left
      self.right = right
class Solution:
   def solve(self, root):
      count = [0, 0, 0]
   def dfs(node):
      if node:
         pre = count[:]
         dfs(node.left)
         dfs(node.right)
         count[node.val] += 1
         node.count = [count[i] - pre[i] for i in range(2)]
   dfs(root)
   def dfs2(node, par=None):
      if node:
         if par is not None:
            a0, a1 = node.count
            b0, b1 = count[0] - a0, count[1] - a1
            if (a0 == 0 or a1 == 0) and (b0 == 0 or b1 == 0):
               self.ans += 1
         dfs2(node.left, node)
         dfs2(node.right, node)
   self.ans = 0
   dfs2(root)
   return self.ans
ob = Solution()
root = TreeNode(0)
root.left = TreeNode(0)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
print(ob.solve(root))

입력

root = TreeNode(0)
root.left = TreeNode(0)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)

출력

1

동작 원리 정리

첫 번째 DFS(dfs)는 후위 순회(post-order traversal) 방식으로 트리를 탐색하면서, 각 노드의 서브트리에 포함된 0과 1의 개수를 미리 계산해 저장합니다. 두 번째 DFS(dfs2)는 간선을 하나 제거했을 때 생기는 두 트리를 시뮬레이션합니다. 한쪽 트리의 구성은 node.count에서, 다른 쪽 트리의 구성은 전체 카운트에서 서브트리 카운트를 뺀 값으로 구할 수 있습니다. 두 트리 중 하나라도 0과 1을 동시에 포함한다면 그 간선은 유효하지 않으며, 양쪽 모두 조건을 만족할 때만 답을 하나씩 늘려갑니다. 이 알고리즘의 시간 복잡도는 O(n), 공간 복잡도 역시 O(n)으로 트리 순회 두 번으로 효율적으로 해결됩니다.