Pythonで誤ったリンクを持つ二分木を検出・修正する方法

誤った二分木とは？

ここでは、ある種の欠陥を持つ二分木を扱います。具体的には、あるノードの右子ポインタが、同じ階層にある別のノードを誤って指してしまっている状態です。この問題を修正するには、誤ったポインタを持つノードを特定し、そのノードとその子孫を木から取り除きます。ただし、誤って指されていたノード自体は木に残します。最後に、修正後の二分木のルートノードを返します。

例として、次のような木が与えられた場合を考えてみましょう。

図のように、ノード4とノード6の間に不正なリンクが張られており、ノード4の右子ポインタがノード6を指しています。

この場合、修正後の木を中順走査（inorder traversal）で表現すると、出力は次のようになります。

2, 3, 5, 6, 7, 8,

ノード4はノード6への誤ったリンクを持っているため、木から削除されます。

解決のためのアルゴリズム

この問題は、幅優先探索（BFS）と訪問済みノードの管理を組み合わせることで解くことができます。手順は以下の通りです。

• ルートを格納した新しいキュー（deque）q を用意します。
• 各ノードの親情報を記録するためのマップ p と、訪問済みノードを管理する集合 visited を用意します。
• キュー q が空になるまで、以下を繰り返します。
  • キューの先頭から要素 cur を取り出します。
  • cur がすでに visited に存在する場合、それは誤ったリンクを通じて再訪問されたことを意味します。
    • p[cur] から cur の親ノードとその位置（左子か右子か）を取得し、さらにその親の情報 p[親] から祖父母ノードと位置を取得します。
    • 親が祖父母の左子であれば祖父母の左ポインタを None に、そうでなければ右ポインタを None に設定します。これにより、誤ったリンクを持つノードが木から切り離されます。
    • ルートを返して処理を終了します。
  • cur を visited に追加します。
  • cur の左子が存在すれば、p[左子] = (cur, 1) と記録し、キューの末尾に追加します。
  • cur の右子が存在すれば、p[右子] = (cur, 0) と記録し、キューの末尾に追加します。
• すべての処理が完了したら、ルートを返します。

Pythonでの実装例

それでは、理解を深めるために実際の実装を見てみましょう。

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 solve(root):
   q = collections.deque([root])
   p, visited = dict(), set()
   while q:
      cur = q.popleft()
      if cur in visited:
         grand_p, is_left = p[p[cur][0]]
         if is_left:
            grand_p.left = None
         else:
            grand_p.right = None
         return root
      visited.add(cur)
      if cur.left:
         p[cur.left] = (cur, 1)
         q.append(cur.left)
      if cur.right:
         p[cur.right] = (cur, 0)
         q.append(cur.right)
   return root
root = make_tree([5, 3, 7, 2, 4, 6, 8])
link_from = search_node(root, 4)
link_to = search_node(root, 6)
link_from.right = link_to
print_tree(solve(root))

入力

root = make_tree([5, 3, 7, 2, 4, 6, 8])
link_from = search_node(root, 4)
link_to = search_node(root, 6)
link_from.right = link_to

出力

2, 3, 5, 6, 7, 8,

アルゴリズムのポイント

この手法の鍵となるのは、visited 集合による重複検出です。正常な二分木であれば、幅優先探索中に同じノードが2回現れることはありません。しかし、誤ったリンクが存在すると、すでに訪問済みのノードが再びキューに追加されるため、そこで異常を検出できます。

また、辞書 p に各ノードの親と左右の位置を事前に記録しておくことで、誤りのあるノードをその親から効率的に切り離せる点も重要です。計算量はノード数を n とすると O(n)、必要な追加メモリも O(n) であり、大きな木に対しても実用的なアプローチと言えます。