Pythonで二分木を前順走査して文字列を構築する方法

二分木が与えられたとき、前順走査（先行順トラバーサル）の方法で木をたどり、括弧と整数からなる文字列を構築することを考えます。ヌルノードは空の括弧のペア「()」で表現します。ただし、文字列と元の二分木との一対一の対応関係に影響しない空の括弧のペアは、すべて省略する必要があります。

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

この場合、出力は 5(6()(8))(7) となります。左の子が存在し、そのさらに右に子があるため「6()(8)」のように空の括弧が必要になりますが、それ以外の不要な空の括弧は省略されています。

解法のアプローチ

この問題を解くために、以下の手順に従います。

• 結果を格納するための空文字列 ans を用意します。
• 再帰関数 pot() を定義します。引数としてノードを受け取ります。
• ノードが null の場合は、空文字列を返します。
• ノードが葉（左右の子がどちらも null）の場合は、そのノードの値を文字列にして返します。
• それ以外の場合は、まず ss にノードの値を設定します。
• 左の子が存在するなら、ss に「(」+ pot(左の子) +「)」を連結します。
• 左の子が存在しない場合は、ss に「()」を連結します。これは右側の構造を正しく表現するために必要です。
• 右の子が存在するなら、ss に「(」+ pot(右の子) +「)」を連結します。右の子が null の場合、対応関係に影響しないため何も追加しません。
• 最後に ss を返します。
• メイン処理では、根ノードに対して pot() を呼び出した結果を返します。

実装例

以下の実装を見ると、理解がより深まるでしょう。

Pythonコード

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

class Solution:
    def tree2str(self, t: TreeNode):
        ans = ''
        def pot(node, s):
            if node == None or node.data == 0:
                return ''
            if node.left == node.right == None:
                return str(node.data)
            ss = str(node.data)
            if node.left != None:
                ss = ss + '(' + pot(node.left, '') + ')'
            else:
                ss = ss + '()'
            if node.right != None:
                ss = ss + '(' + pot(node.right, '') + ')'
            return ss
        return pot(t, '')

ob = Solution()
root = make_tree([5,6,7,None,8])
print(ob.tree2str(root))

入力

[5,6,7,None,8]

出力

5(6()(8))(7)

まとめ

このアルゴリズムは、二分木を前順走査しながら再帰的に文字列を組み立てていくシンプルな手法です。ポイントは、左の子が存在しない場合に空の括弧「()」を残す必要がある一方で、右の子が存在しない場合は括弧を省略できるという点です。これにより、文字列から元の二分木を一意に復元できる形式が保たれます。計算量は各ノードを一度だけ訪問するため、ノード数を n とすると時間計算量 O(n)、再帰の深さによる空間計算量も最悪で O(n) となります。