โปรแกรมนับจำนวนวิธีแบ่งต้นไม้เป็นสองต้นไม้ใน Python

สมมติว่าเรามีต้นไม้ทวิภาคมีค่า 0, 1, 2 รากมีโหนด 0 อย่างน้อย 1 ตัวและโหนด 1 อย่างน้อย 1 ตัว มีการลบเส้นเชื่อม (edge) ทำให้กลายเป็นสองต้นไม้ เราต้องหาจำนวนวิธีลบเส้นเชื่อม 1 เส้น ที่ทำให้ทั้งสองต้นไม้ไม่มีทั้งโหนด 0 และ 1 ร่วมกัน

ตัวอย่าง: จากภาพ ผลลัพธ์คือ 1 (ลบเส้นเชื่อม 0-2 เท่านั้น)

ขั้นตอนการแก้ปัญหา

1. count = [0, 0, 0]
2. ฟังก์ชัน dfs(node):
   • ถ้า node ไม่ใช่ None:
     • pre = count[:]
     • เรียก dfs(node.left), dfs(node.right)
     • count[node.val] += 1
     • node.count = [count[i] - pre[i] for i in range(2)]
3. ฟังก์ชัน dfs2(node, par=None):
   • ถ้า node ไม่ใช่ None:
     • ถ้า par ไม่ใช่ None:
       • (a0, a1) = node.count
       • (b0, b1) = (count[0] - a0, count[1] - a1)
       • ถ้า (a0 == 0 or a1 == 0) and (b0 == 0 or b1 == 0): ans += 1
     • เรียก dfs2(node.left, node), dfs2(node.right, node)
4. จากเมธอดหลัก: เรียก dfs(root), ans = 0, เรียก dfs2(root), คืนค่า ans

ตัวอย่างโค้ด Python

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

อินพุต / เอาต์พุต

Input: ต้นไม้ตามโค้ดตัวอย่าง

Output: 1