Chương trình đếm số cách chia cây nhị phân thành hai cây bằng Python

Cho cây nhị phân chứa giá trị 0, 1, 2. Gốc có ít nhất một nút 0 và một nút 1. Xóa một cạnh để tách thành hai cây. Đếm số cách xóa sao cho không cây nào chứa cả 0 và 1.

Ví dụ minh họa cho kết quả là 1 (chỉ xóa cạnh 0-2).

Quy trình giải quyết

1. DFS lần 1: tính số lượng nút 0 và 1 trong cây con của mỗi nút, lưu vào node.count.
2. DFS lần 2: từ gốc, xét mỗi cạnh cha-con. Nếu tách cạnh đó, cây con có (a0, a1), cây còn lại có (b0, b1). Nếu (a0==0 or a1==0) and (b0==0 or b1==0) thì tăng đáp án.
3. Trả về tổng số cách.

Ví dụ mã nguồn

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):
        count = [0, 0, 0]
        def dfs(node):
            if node:
                pre = count[:]
                dfs(node.left)
                dfs(node.right)
                count[node.val] += 1
                node.count = [count[i] - pre[i] for i in range(2)]
        dfs(root)
        self.ans = 0
        def dfs2(node, par=None):
            if node:
                if par is not None:
                    a0, a1 = node.count
                    b0, b1 = count[0] - a0, count[1] - a1
                    if (a0 == 0 or a1 == 0) and (b0 == 0 or b1 == 0):
                        self.ans += 1
                dfs2(node.left, node)
                dfs2(node.right, node)
        dfs2(root)
        return self.ans

ob = Solution()
root = TreeNode(0)
root.left = TreeNode(0)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
print(ob.solve(root))  # Output: 1

Đầu vào

root = TreeNode(0) root.left = TreeNode(0) root.right = TreeNode(2) root.right.left = TreeNode(1) root.right.right = TreeNode(1)

Đầu ra

1