Pythonで二分木の最大平均サブツリーを求めるアルゴリズム

問題の概要

二分木のルートが与えられたとき、その木に含まれる任意の部分木の平均値の最大値を求めるのがこの問題の目的です。例として、次のような二分木を考えてみましょう。

この場合の出力は 6 になります。理由は以下の通りです。

• ノード5を根とする部分木：(5 + 6 + 1) / 3 = 4
• ノード6を根とする部分木：6 / 1 = 6
• ノード1を根とする部分木：1 / 1 = 1

これらの中で最も大きいのは6なので、答えは6となります。

解法のアプローチ

この問題は後順走査（post-order traversal）による再帰を使うことで効率的に解けます。各ノードについて「その部分木のノード数」と「その部分木の合計値」をペアで返しながら、途中経過で平均値を更新していくのがポイントです。

具体的には、以下の手順で処理を進めます。

1. 結果を格納する変数 res を0で初期化します。
2. solve() というメソッドを定義し、ノードを引数として受け取ります。
3. ノードが存在しない場合は、ペア [0, 0]（ノード数0、合計値0）を返します。
4. 左右の子ノードに対して solve() を再帰的に呼び出します。
5. 現在の部分木のノード数 c を「左のノード数 + 右のノード数 + 1」として計算します。
6. 現在の部分木の合計値 s を「左の合計値 + 右の合計値 + 現在のノードの値」として計算します。
7. ans を「ans と s / c のうち大きい方」で更新します。
8. ペア [c, s] を呼び出し元に返します。
9. メインのメソッドでは ans を0に初期化し、solve(root) を呼び出した後に ans を返します。

実装例（Python）

理解を深めるために、以下の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(object):
    def helper(self, node):
        if not node:
            return 0, 0
        left_sum, left_count = self.helper(node.left)
        right_sum, right_count = self.helper(node.right)
        self.ans = max(self.ans, (left_sum + right_sum + node.data) / (left_count + right_count + 1))
        return left_sum + right_sum + node.data, left_count + right_count + 1

    def maximumAverageSubtree(self, root):
        self.ans = 0
        self.helper(root)
        return self.ans

ob = Solution()
root = make_tree([5, 6, 1])
print(ob.maximumAverageSubtree(root))

入力

[5,6,1]

出力

6.0

計算量の評価

このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量はO(n)（nはノード数）です。また、再帰の深さは木の高さに依存するため、バランスの取れた木では空間計算量はO(log n)、最悪ケース（線形リスト状に偏った木）ではO(n)となります。

なお、ノードの値に負の数が含まれる可能性がある場合は、ans の初期値を0ではなく float('-inf')（負の無限大）にしておくと、より安全に動作します。