Pythonで二分木の上面図（トップビュー）を求めるプログラムの書き方

二分木の上面図とは

二分木が与えられたとき、その木を真上から見たときに見えるノードの値を「上面図（トップビュー）」と呼びます。結果は必ず左から右の順に並べて出力します。

例えば、次のような入力があった場合、出力は [3, 5, 8, 6, 9] となります。これは、ノード3がノード2の真上に、ノード5がノード7の真上に位置しているため、2と7は上から見ると隠れてしまい、見えないからです。

アルゴリズムの考え方

この問題は、幅優先探索（BFS）と各ノードの「水平座標」を使うことで効率的に解くことができます。ルートの座標を0とし、左の子へ移動するたびに座標を-1、右の子へ移動するたびに座標を+1として管理します。同じ座標を持つノードの中では、BFSの性質上、最初に到達したノード（＝最も上にあるノード）だけが見えることになります。

具体的な手順は以下の通りです。

• view := 空のマップ（辞書）を用意する
• q := 両端キュー（deque）を用意する
• (root, 0) のペアを q の末尾に追加する
• start := 無限大、end := 負の無限大で初期化する
• q が空になるまで以下を繰り返す
  • (node, coord) := q の先頭要素を取り出す
  • start := start と coord の最小値
  • end := end と coord の最大値
  • coord が view に存在しない場合は、view[coord] := ノードの値 を記録する
  • node の左の子が存在するなら、(左の子, coord - 1) を q に追加する
  • node の右の子が存在するなら、(右の子, coord + 1) を q に追加する
• res := 新しいリストを作成する
• i を start から end まで動かしながら、i が view に存在すれば view[i] を res の末尾に追加する
• res を返す

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

Pythonコード例

from collections import deque
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):
        view = {}
        q = deque()
        q.append((root, 0))
        start = float("inf")
        end = float("-inf")
        while q:
            node, coord = q.popleft()
            start = min(start, coord)
            end = max(end, coord)
            if coord not in view:
                view[coord] = node.val
            if node.left:
                q.append((node.left, coord - 1))
            if node.right:
                q.append((node.right, coord + 1))
        res = []
        for i in range(start, end + 1):
            if i in view:
                res.append(view[i])
        return res
ob = Solution()
root = TreeNode(5)
root.left = TreeNode(3)
root.right = TreeNode(8)
root.right.left = TreeNode(7)
root.right.right = TreeNode(6)
root.right.left.left = TreeNode(2)
root.right.right.right = TreeNode(9)
print(ob.solve(root))

入力

root = TreeNode(5) root.left = TreeNode(3) root.right = TreeNode(8)
root.right.left = TreeNode(7) root.right.right = TreeNode(6)
root.right.left.left = TreeNode(2) root.right.right.right =
TreeNode(9)

出力

[3, 5, 8, 6, 9]

計算量について

このアルゴリズムの時間計算量は O(n log n)、空間計算量は O(n) となります。ここで n は二分木のノード数です。BFSですべてのノードを一度ずつ訪問し、最後に座標の範囲でソートされた結果を組み立てるため、非常に効率的な手法といえます。