Chương trình tìm số đường đi độ dài k trên cây nhị phân trong Python

Cho cây nhị phân giá trị duy nhất và số k, đếm số đường đi độ dài k (đi từ cha đến con hoặc con đến cha). Hai đường đi khác nhau nếu có nút khác nhau.

Ví dụ

Cây đầu vào: (hình ảnh), k = 3.
Đầu ra: 4 (đường đi: [12,8,3], [12,8,10], [8,12,15], [3,8,10]).

Thuật toán

1. DFS trả về mảng res kích thước K, res[i] = số đường đi từ nút hiện tại xuống con cháu ở độ sâu i.
2. Kết hợp kết quả từ trái và phải: ans += left[i] * right[K-1-i].
3. Cập nhật res cho nút cha.
4. Trả về ans.

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, K):
        def dfs(node):
            if not node:
                return [1] + [0] * (K-1)
            left = dfs(node.left)
            right = dfs(node.right)
            for i in range(K):
                self.ans += left[i] * right[K - 1 - i]
            res = [0] * K
            res[0] = res[1] = 1
            for i in range(1, K - 1):
                res[i + 1] += left[i]
                res[i + 1] += right[i]
            return res
        self.ans = 0
        dfs(root)
        return self.ans

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

Đầu vào / Đầu ra

Input: root = TreeNode(12) ..., 3
Output: 4