Pythonで二分木からノードを削除し、残りのフォレストの根を求める方法

本記事では、二分木から特定のノードを削除した際に生じる「フォレスト（森）」の根を求めるアルゴリズムを、Pythonのコードとともにわかりやすく解説します。

問題の概要

二分木の根（root）が与えられ、木に含まれる各ノードは一意の値を持っているものとします。ここで、to_delete リストに含まれる値を持つノードをすべて削除すると、木はいくつかの独立した部分木、すなわち「フォレスト」へと分割されます。このとき、残ったフォレストを構成する各木の根を見つけるのが目的です。

たとえば、次のような二分木が入力として与えられたとします。

このとき to_delete 配列が [3, 5] である場合、値3と5のノードが削除され、出力は次のようになります。

解法のアプローチ

この問題は、深さ優先探索（DFS）による再帰処理を使うことで効率的に解けます。手順は以下の通りです。

• 結果を格納するための配列 res を定義します。
• 再帰用のメソッド solve() を定義します。引数は「現在のノード」「to_delete 配列」「そのノードが根であるかどうかを示すブール値 is_root」です。
• ノードが null の場合は null を返します。
• フラグ flag を、ノードの値が to_delete に含まれている場合に true となるよう設定します。
• flag が false かつ is_root が true の場合、そのノードを res に追加します。
• 左の子に対して solve(左の子, to_delete, flag) を呼び出します。
• 右の子に対して solve(右の子, to_delete, flag) を呼び出します。
• flag が立っている場合は None を返し、そうでなければノード自身を返します。
• メインメソッドからは solve(root, to_delete, True) の形で呼び出します。

ポイントは、親ノードが削除対象だった場合、その子ノードは新たなフォレストの根の候補になるという点です。flag を子ノードへの is_root 引数として渡すことで、この関係を自然に表現できます。

実装例

class Solution(object):
    def delNodes(self, root, to_delete):
        """
        :type root: TreeNode
        :type to_delete: List[int]
        :rtype: List[TreeNode]
        """
        to_delete = set(to_delete)
        self.res = []
        self.solve(root, to_delete, True)
        return self.res

    def solve(self, node, to_delete, is_root):
        if not node:
            return None
        flag = node.val in to_delete
        if not flag and is_root:
            self.res.append(node)
        node.left = self.solve(node.left, to_delete, flag)
        node.right = self.solve(node.right, to_delete, flag)
        return None if flag else node

コードのポイント

to_delete をリストから set に変換している点に注目してください。Pythonのセットは要素の存在確認が平均O(1)で行えるため、各ノードごとの判定が高速になります。

また、このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量はO(n)、空間計算量も再帰の深さにより最悪O(n)となります。

動作確認

入力

[1,2,3,4,5,6,7]
[3,5]

出力

[[1,2,null,4],[6],[7]]

値3と5のノードが削除され、残った3つの木の根 [1,2,null,4]、[6]、[7] が正しく返されていることが確認できます。