LeetCode の Delete Nodes And Return Forest を Python 3 で解いてみた。

アルゴリズム

僕は木を探索しながら返り値を作り上げていくような実装をした。

具体的には、再帰関数 (walk_tree) に木の頂点を渡し、そこから子ノードがある限り自分自身をコールしていく。

再帰関数には has_parent という引数を設けて、親ノードがある場合とない場合で処理を分岐させる。親ノードがない場合 (i.e., has_parent = False の場合) はルートノードなので、返り値に自分自身を追加する。

木をたどる中で、ノードの値が to_delete に含まれる場合は has_parent = False として再帰関数をコールする。また 親を持たないノードはルート なので、そのノードは返り値のリストに append する。

コード

具体的なコードは次の通り。この実装で Runtime: 68 ms, faster than 66.80% of Python3 online submissions for Delete Nodes And Return Forest. とのこと。なお、同じ実装で faster than 85.44% of Python3 online submissions になることもあるので、多くの そこそこ速い 実装は似ているのだと思う。

ブログの読者向けに詳細にコメントを書き込んだコードを貼り付けておく。

from typing import List


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


class Solution:
    def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:
        # メソッド内で ans を組み立てて返り値として返却する
        ans = []

        # "root.val が to_delete に含まれるのか" を複数回調べたいので、高速化のため
        # to_delete を set にしておく
        # 補足: `something in LIST` は O(N) で `something in SET` は O(1)
        to_delete_set = set(to_delete)

        # walk_tree() は再帰関数 (node の子ノードを引数にしてコールする)
        def walk_tree(node, has_parent):
            # None までたどり着いたら None を返却する
            # 葉に到達した場合か to_delete に含まれるノードにきた場合に `node is None` になる
            if node is None:
                return None

            if node.val in to_delete_set:
                # node は削除対象なので次のふたつを行う
                #
                #  1. `has_parent = False` で再帰する
                #  2. None を返却する
                node.left = walk_tree(node.left, False)
                node.right = walk_tree(node.right, False)

                return None
            else:
                if not has_parent:
                    # 親がいない場合 (i.e., ルートノードである場合) は自身を ans に追加する
                    ans.append(node)

                # node.val が削除対象ではないので `has_parent = True` で再帰する
                node.left = walk_tree(node.left, True)
                node.right = walk_tree(node.right, True)

                # 自分自身を返す (これを忘れると None が返ってしまう)
                return node

        # root ノードに親はいないため、最初は has_parent = False でコールする
        walk_tree(root, False)

        return ans