Pythonで二分木の対角線ごとの要素の合計を求める方法

問題の概要

二分木が与えられたとき、木の各対角線（右上から左下に向かう経路）ごとに、その対角線上のノード値の合計を求めることを考えます。

たとえば、次のような二分木が入力だったとします。

この場合、対角線は [12, 15]、[8, 10]、[3] の3本になるため、それぞれの合計を求めると出力は [27, 18, 3] となります。

アルゴリズム

この問題を解くには、次の手順に従います。まず traverse() 関数を定義します。この関数は node、numLeft、output の3つの引数を受け取ります。

• node が null（None）の場合は、そのまま return します。
• numLeft が output のサイズ以上の場合は、node のデータを output の末尾に追加します。
• それ以外の場合は、output[numLeft] += node.data として該当する対角線の合計を更新します。
• node の左の子が存在する場合は、traverse(node.left, numLeft + 1, output) を呼び出します。
• node の右の子が存在する場合は、traverse(node.right, numLeft, output) を呼び出します。

メインメソッドでは、以下の処理を行います。

• 結果を格納する新しいリスト output を作成します。
• traverse(root, 0, output) を呼び出して探索を開始します。
• output を返します。

ポイント解説

このアルゴリズムの鍵となるのは、「右方向への移動では同じ対角線にとどまり、左方向への移動で次の対角線に移る」という二分木の性質です。そのため、右の子へ再帰するときは numLeft をそのまま渡し、左の子へ再帰するときは numLeft + 1 を渡すことで、各ノードが属する対角線のインデックスを正しく管理できます。

実装例

理解を深めるために、以下のPythonによる実装例を見てみましょう。

class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.data = data
        self.left = left
        self.right = right
class Solution:
    def solve(self, root):
        output = []
        def traverse(node, numLeft, output):
            if not node:
                return
            if numLeft >= len(output):
                output.append(node.data)
            else:
                output[numLeft] += node.data
            if node.left:
                traverse(node.left, numLeft+1, output)
            if node.right:
                traverse(node.right, numLeft, output)
        traverse(root, 0, output)
        return output
ob = Solution()
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
print(ob.solve(root))

入力

root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)

出力

[27, 18, 3]