Pythonで二分木の指定ノードの右隣ノードを見つけるプログラム

二分木が与えられ、さらに特定のノード「u」へのポインタも渡されたとします。このとき、u のすぐ右側に位置するノード(必ず同じ階層に存在する)を見つける必要があります。対象のノードは葉ノードの場合もあれば、内部ノードの場合もあります。

例として、次のような二分木が入力されたとしましょう。

ここで u = 6 とすると、出力は 8 になります。ノード 6 の右隣にはノード 8 が存在するため、値 8 が返されるというわけです。

解決のためのアプローチ

この問題は、両端キュー(deque)を使った幅優先探索(BFS)、いわゆるレベル順走査によって解くことができます。手順は以下の通りです。

• ルートが空であれば、null を返す
• dq := 新しい deque を作成する
• dq の末尾にルートを挿入する
• dq が空でない間、以下を繰り返す
  • dq_size := dq のサイズ
  • temp := 新しいリスト、index := -1
  • 0 から dq_size までの範囲で各値に対して
    • node := dq の末尾から要素を取り出す
    • node の左の子が存在すれば、dq の末尾に追加する
    • node の右の子が存在すれば、dq の末尾に追加する
    • node を temp の末尾に挿入する
    • node が u と同一であれば、index := len(temp) - 1 とする
  • index が temp の最後尾と一致する場合、null を返す(右隣のノードが存在しない)
  • index > -1 であれば、temp[index + 1] を返す
• null を返す

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

実装例(Python)

from queue import deque
class TreeNode:
   def __init__(self, val=0, left=None, right=None):
      self.val = val
      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.val == element):
      return root
   res1 = search_node(root.left, element)
   if res1:
      return res1
   res2 = search_node(root.right, element)
   return res2
def solve(root, u):
   if not root:
      return None
   dq = deque()
   dq.append(root)
   while dq:
      dq_size = len(dq)
      temp = []
      index = -1
      for _ in range(dq_size):
         node = dq.pop()
         if node.left:
            dq.appendleft(node.left)
         if node.right:
            dq.appendleft(node.right)
         temp.append(node)
         if node == u:
            index = len(temp) - 1
      if index == len(temp) - 1:
         return None
      if index > -1:
         return temp[index + 1]
   return None
root = make_tree([5, 3, 7, 2, 4, 6, 8])
u = search_node(root,6)
ret = solve(root, u)
print(ret.val)

入力

root = make_tree([5, 3, 7, 2, 4, 6, 8])
u = search_node(root,6)

出力

8