Pythonで二分木の合計がkとなるパスの数を数える方法

問題の概要

二分木と値 k が与えられたとき、あるノードからその子孫へ向かうパスのうち、通過するノードの値の合計がちょうど k と一致するものがいくつ存在するかを求める問題です。

例えば、次のような二分木を考えてみましょう。

このとき k = 5 であれば、出力は 2 となります。条件を満たすパスは [2, 3] と [1, 4] の2つだからです。

解き方のアプローチ:累積和(prefix sum)の活用

この問題は「累積和(prefix sum)」というテクニックを使うことで、全ノードを一度だけ訪問する効率的なアルゴリズムとして解けます。考え方の手順は以下の通りです。

• count:マップ(Counter)を用意し、キー 0 に対して初期値 1 を設定しておきます
• ans(答え)と prefix(ルートから現在ノードまでの累積和)を 0 で初期化します
• dfs() 関数を定義します。引数はノードです
• ノードが null でない場合、以下の処理を行います:
  • prefix に現在のノードの値を加算する
  • ans に count[prefix - target] を加算する(キーが存在しない場合は 0 として扱う)
  • count[prefix] を +1 する
  • 左の子ノードに対して dfs を再帰呼び出しする
  • 右の子ノードに対して dfs を再帰呼び出しする
  • count[prefix] を -1 して元に戻す(バックトラック)
  • prefix からノードの値を減算して元に戻す(バックトラック)
• メイン処理では dfs(root) を呼び出し、最後に ans を返します

この手法のポイントは、「現在位置までの累積和から k を引いた値」が過去に何回現れたかを記録しておくことで、途中の任意の区間の合計が k に一致するかどうかを即座に判定できる点にあります。また、再帰から戻る際に count と prefix を元に戻すことで、別の枝のパス同士が互いに干渉しないようになっています。

実装例

それでは、実際のPythonコードを見てみましょう。

from collections import Counter

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, target):
        count = Counter([0])
        ans = prefix = 0

        def dfs(node):
            nonlocal ans, prefix
            if node:
                prefix += node.val
                ans += count[prefix - target]
                count[prefix] += 1
                dfs(node.left)
                dfs(node.right)

                # バックトラック:状態を元に戻す
                count[prefix] -= 1
                prefix -= node.val

        dfs(root)
        return ans


ob = Solution()
root = TreeNode(3)
root.left = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(1)
root.right.left.right = TreeNode(2)
k = 5
print(ob.solve(root, k))

入力

root = TreeNode(3)
root.left = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(1)
root.right.left.right = TreeNode(2)
k = 5

出力

2

計算量について

このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量は O(n) です。また、マップに保存される累積和のエントリ数は最大でもノード数に比例するため、空間計算量も O(n) となります。素朴な全パス列挙のアプローチ(O(n²)程度かかる場合がある)と比べて大幅に効率化されているのがわかります。