Chương trình tính tổng các phần tử trên từng đường chéo của cây nhị phân trong Python

Cho cây nhị phân, tính tổng của mỗi đường chéo từ trên xuống dưới bên phải.

Ví dụ

Cây đầu vào: (hình ảnh).
Đầu ra: [27, 18, 3] (đường chéo: [12,15], [8,10], [3]).

Thuật toán

1. Định nghĩa hàm traverse(node, numLeft, output):
   • Nếu node là None, trả về.
   • Nếu numLeft >= len(output), thêm node.data vào output.
   • Ngược lại, cộng node.data vào output[numLeft].
   • Đệ quy trái với numLeft+1, phải với numLeft.
2. Khởi tạo output = [], gọi traverse(root, 0, output).
3. Trả về output.

Mã nguồn 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))

Đầu vào / Đầu ra

Input: root = TreeNode(12) ...
Output: [27, 18, 3]