Pythonで二分木を2つの木に分割できるパターン数を数えるプログラム

問題の概要

値「0」「1」「2」を含む二分木があるとします。根（ルート）には、少なくとも1つの「0」ノードと1つの「1」ノードが存在しています。ここで、「木の辺（エッジ）を1本削除すると、木が2つの異なる木に分割される」という操作を考えます。

このとき、削除後に生成される2つの木のどちらにも「0」と「1」のノードが同時に含まれないように、辺を1本削除する方法が何通りあるかを求めるのがこの問題です。

入力例

例えば、次のような二分木が与えられたとします。

この場合、出力は 1 となります。「0」から「2」へ向かう辺だけが、条件を満たす唯一の削除対象だからです。

解法のアプローチ

この問題は、DFS（深さ優先探索）を2回実行することで効率的に解けます。

1. 1回目のDFS（dfs）：各ノードについて、そのノードを根とする部分木に含まれる「0」と「1」の個数を事前に計算しておきます。
2. 2回目のDFS（dfs2）：各辺（親子間の接続）について、その辺を削除した場合に両側の部分木が条件を満たすかどうかを判定し、条件を満たす辺の数をカウントします。

アルゴリズムの手順

• カウンター count を [0, 0, 0] で初期化します。
• 関数 dfs(node) を定義します。
  • node が null でない場合：
    • pre に現在の count のコピーを保存します。
    • 左の子・右の子に対して再帰的に dfs を呼び出します。
    • count[node.val] を1増やします。
    • node.count に、この部分木内の「0」と「1」の個数（count[i] - pre[i]、i = 0, 1）を記録します。
• 関数 dfs2(node, par) を定義します。
  • node が null でない場合：
    • par が null でない場合（親との辺が存在する場合）：
      • (a0, a1) := node.count（子側の部分木の「0」「1」の個数）
      • (b0, b1) := (count[0] - a0, count[1] - a1)（親側の部分木の「0」「1」の個数）
      • (a0 == 0 または a1 == 0) かつ (b0 == 0 または b1 == 0) を満たす場合、ans を1増やします。
    • 左の子・右の子に対して再帰的に dfs2 を呼び出します。
• メイン処理では以下を実行します：
  • dfs(root) を呼び出して各部分木のカウントを計算します。
  • ans := 0 で初期化します。
  • dfs2(root) を呼び出して条件を満たす辺の数を数えます。
  • ans を返します。

実装例（Python）

理解を深めるために、以下の実装を見てみましょう。

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

計算量とまとめ

このアルゴリズムでは、各ノードを2回のDFSでそれぞれ1度ずつ訪問するため、時間計算量は O(N) となります。また、再帰の呼び出しスタックと各ノードのカウント保持にO(N) の空間が必要です。

ポイントは、1回目のDFSで「各部分木ごとの0と1の個数」を前計算しておくことで、2回目のDFSで各辺の判定を O(1) で行える点です。これにより、全ての辺を効率的に評価できます。