Pythonで解く!二分木のルートからリーフへのパスにおける不十分なノードの削除方法

問題概要

二分木が与えられたとき、あるノードが「不十分(insufficient)」であるとは、そのノードを通るすべてのルートからリーフへのパスのノード値の合計が、与えられた limit よりも厳密に小さいことを意味します。この条件を満たすすべての不十分なノードを同時に削除し、処理後の二分木のルートを返すのが本問題の目的です。

例えば、次のような二分木があり、limit = 1 が与えられたとします。

このとき、不十分なノードを削除した後の出力は以下のようになります。

解法のアプローチ

この問題は、再帰(深さ優先探索)を使って効率的に解くことができます。基本的な考え方は、各リーフノードについてパスの合計が limit を満たしているかを判定し、条件を満たさない部分木を順に切り離していくというものです。

アルゴリズムの手順

• solve(root, limit) メソッドを定義する。
• ノードがリーフ(左右の子を持たない)の場合:
  • ノードの値が limit 未満であれば null(None)を返し、そうでなければそのノード自身を返す。
• 左の子が存在する場合は、root.left = solve(root.left, limit - root.val) で再帰的に処理する。
• 右の子が存在する場合は、root.right = solve(root.right, limit - root.val) で再帰的に処理する。
• 最後に、左右どちらかの子が残っていればルートを返し、両方とも削除された場合は null(None)を返す。

Pythonでの実装例

class Solution(object):
    def sufficientSubset(self, root, limit):
        """
        :type root: TreeNode
        :type limit: int
        :rtype: TreeNode
        """
        if not root.left and not root.right:
            return None if root.val < limit else root
        if root.left:
            root.left = self.sufficientSubset(root.left, limit - root.val)
        if root.right:
            root.right = self.sufficientSubset(root.right, limit - root.val)
        return root if root.left or root.right else None

実行例

入力:

[1,2,3,4,-99,-99,7,8,9,-99,-99,12,13,-99,14]
1

出力:

[1,2,3,4,null,null,7,8,9,null,14]

計算量の分析

このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量は O(n)(n はノード数)です。また、再帰呼び出しの深さは木の高さに依存するため、空間計算量は最悪ケースで O(n)、バランスの取れた二分木の場合は O(log n) となります。

まとめ

本記事では、二分木のルートからリーフへのパスの合計が limit 未満となる「不十分な」ノードを同時に削除する問題を、再帰的な深さ優先探索を用いて解く方法を紹介しました。リーフノードを基準に判定を行い、子の状態に応じて親ノードも削除するかどうかを決定する、シンプルながら強力なアプローチです。類似の木構造の問題にも応用できるテクニックなので、ぜひマスターしておきましょう。