Chương trình đếm số đường đi trong cây nhị phân có tổng bằng k bằng Python

Cho cây nhị phân và giá trị k. Đếm số đường đi từ nút cha đến nút con (xuống theo cây) có tổng giá trị bằng k.

Ví dụ: cây minh họa, k=5 → 2 đường đi ([2,3] và [1,4]).

Giải thuật

• Sử dụng DFS duyệt tiền tự (preorder) kết hợp tiền tố (prefix sum) và bảng băm count lưu tần suất các tổng tiền tố.
• Khởi tạo count[0] = 1, ans = 0, prefix = 0.
• Tại mỗi nút: prefix += node.val, ans += count[prefix - target], count[prefix] += 1.
• Đệ quy con trái, con phải.
• Khi backtrack: count[prefix] -= 1, prefix -= node.val.

Ví dụ minh họa

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))

Đầu vào

root = TreeNode(3) ... k = 5

Đầu ra

2