Pythonで二分木の中に連結リストと一致するパスが存在するか判定する方法

問題の概要

根ノード「root」を持つ二分木と、先頭ノード「head」を持つ連結リストが与えられたとします。このとき、連結リストが二分木の中に存在するかどうかを判定します。具体的には、木の中の一連のノードが親から子へと順番につながっており、その並びが与えられた連結リストと完全に一致する場合には「True」を返し、一致しない場合には「False」を返します。

例えば、入力が以下のようなケースを考えてみましょう。

二分木

連結リスト

この場合、二分木の中に 6 → 7 → 10 という並びのパスが存在するため、出力は True になります。

解法のアプローチ

この問題は、文字列検索アルゴリズムとして有名なKMP法（Knuth–Morris–Pratt法）の考え方を木構造の探索に応用することで、効率よく解くことができます。単純な全探索では一致に失敗するたびに先頭からやり直す必要がありますが、「失敗関数（failure function）」を事前に計算しておくことで、途中まで一致していた情報を活かしながら探索を続けられます。

手順は以下の通りです。

• arr := 新しいリストを作成
• size := arr のサイズ
• temp_arr := サイズ (size + 1) の配列を、すべて -1 で初期化
• 関数 helper() を定義する。引数は root、val
  • val >= size ならば True を返す
  • root が None ならば False を返す
  • val := val + 1
  • val > 0 かつ root の値が arr[val - 1] と等しくない間、val := temp_arr[val - 1] + 1 を繰り返す
  • helper(root の左の子, val) または helper(root の右の子, val) が True ならば True を返す
  • それ以外の場合は False を返す
• start := head
• start が null でない間、次を繰り返す
  • arr の末尾に start の値を追加する
  • start := start の次のノード
• node を 1 から size + 1 までループさせる
  • temp_arr[node] := temp_arr[node - 1] + 1
  • temp_arr[node] > 0 かつ arr[node - 1] != arr[temp_arr[node] - 1] の間、temp_arr[node] := temp_arr[temp_arr[node] - 1] + 1 を繰り返す
• helper(root, 0) の結果を返す

実装例

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

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

class ListNode:
   def __init__(self, val, next=None):
      self.val = val
      self.next = next

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):
   node = TreeNode(elements[0])
   for element in elements[1:]:
      insert(node, element)
   return node

def make_list(elements):
   head = ListNode(elements[0])
   for element in elements[1:]:
      ptr = head
      while ptr.next:
         ptr = ptr.next
      ptr.next = ListNode(element)
   return head

def solve(root, head):
   arr = []
   start = head
   while start:
      arr += (start.val,)
      start = start.next
   size = len(arr)
   temp_arr = [-1] * (size + 1)
   for node in range(1, size + 1):
      temp_arr[node] = temp_arr[node - 1] + 1
      while temp_arr[node] > 0 and arr[node - 1] != arr[temp_arr[node] - 1]:
         temp_arr[node] = temp_arr[temp_arr[node] - 1] + 1
   def helper(root, val):
      if val >= size:
         return True
      if not root:
         return False
      val += 1
      while val > 0 and root.val != arr[val - 1]:
         val = temp_arr[val - 1] + 1
      if helper(root.left, val) or helper(root.right, val):
         return True
      return False
   return helper(root, 0)

root = make_tree([6, 7, 8, 9, 10])
head = make_list([6, 7, 10])
print(solve(root, head))

入力

root = make_tree([6, 7, 8, 9, 10])
head = make_list([6, 7, 10])
print(solve(root, head))

出力

True

処理内容の解説

まず、連結リストの各ノードの値を順に取り出して配列 arr に格納し、その配列に対してKMP法と同様の失敗関数 temp_arr を構築します。その後、二分木を深さ優先で探索しながら、現在のノードの値と連結リストの対応する位置の値を比較します。不一致が発生した場合でも、失敗関数を参照することで比較位置を適切に巻き戻せるため、無駄な再探索を省けます。一致した長さがリスト全体のサイズ size に達した時点で、パスが存在すると判断できます。

このアルゴリズムの時間計算量は、二分木のノード数を N、連結リストの長さを M とすると O(N) であり、空間計算量は失敗関数用の配列と再帰スタックの分を合わせて O(M + H)（H は木の高さ）程度になります。バックトラックを単純に行う素朴な手法（最悪 O(N × M)）と比べて大幅に効率化できる点が大きな特徴です。