Pythonでn分木（N-ary Tree）のルートノードを見つけるプログラム

n分木（N-ary Tree）の各ノードが配列として与えられているとします。ここで求めたいのは、木を再構築したうえでルートノードを見つけて返すことです。返されたノードを起点に、木全体を先行順（Preorder）で表示できれば成功です。

たとえば、入力が次のような場合を考えてみましょう。

このときの出力は以下のようになります。

[14, 27, 32, 42, 56, 65]

この出力は、見つけたルートノードから木の先行順走査（Preorder Traversal）を行った結果です。つまり、正しいルートさえ特定できれば、そこから木全体の構造を復元できます。

解決のアプローチ：入次数（In-degree）に注目する

この問題を解く鍵は、ルートノードだけが親を持たないという性質です。すべての子ノードは必ずどこかのノードから参照されているため、「子として参照された回数」＝入次数を数えれば、入次数が0のノードこそがルートだと特定できます。

具体的には、以下の手順で処理を進めます。

• 整数値を格納できる新しいマップ indegree を用意する。

• 木の中の各ノードについて、その子ポインタ内の各子に対して indegree[子の値] を1ずつ増やす。

• 再度木の中の各ノードを調べ、indegree[ノードの値] が0であるノードを見つけたら、それをルートとして返す。

• 該当するノードが存在しなければ null を返す。

Pythonでの実装例

理解を深めるために、実際のコードを見てみましょう。

import collections
class Node:
   def __init__(self, value, child = None) -> None:
      self.val = value
      self.children = []
      if child != None:
         for value in child:
            self.children.append(value)

def solve(tree):
   indegree = collections.defaultdict(int)
   for node in tree:
      for child in node.children:
         indegree[child.val] += 1
   for node in tree:
      if indegree[node.val] == 0:
         return node
   return None

def treeprint(node, tree):
   if node == None:
      tree.append("None")
      return tree
   if tree == None:
      tree = []
   tree.append(node.val)
   for child in node.children:
      treeprint(child, tree)
   return tree

node6 = Node(65)
node5 = Node(56)
node4 = Node(42, [node5, node6])
node3 = Node(32)
node2 = Node(27)
node1 = Node(14, [node2, node3, node4])
tree = [node2, node1, node5, node3, node6, node4]

root = solve(tree)
print(treeprint(root, None))

コードのポイント

• Node クラスは、ノードの値 val と子ノードのリスト children を持ちます。コンストラクタに子ノードのリストを渡すことで、親子関係を構築できます。

• solve() 関数では、collections.defaultdict(int) を使って入次数マップを作成しています。存在しないキーへのアクセスも自動的に0として扱えるため、コードがシンプルになります。

• treeprint() 関数は再帰的に子ノードをたどりながら値をリストに追加することで、先行順走査の結果を構築します。

入力

node6 = Node(65)
node5 = Node(56)
node4 = Node(42, [node5, node6])
node3 = Node(32)
node2 = Node(27)
node1 = Node(14, [node2, node3, node4])
tree = [node2, node1, node5, node3, node6, node4]

出力

[14, 27, 32, 42, 56, 65]

ノードがバラバラの順序で配列に入っていても、この手法なら確実にルートを特定できます。計算量は全ノードと全エッジを一度ずつ調べるだけなので、O(N) で効率的に動作する点も魅力です。