โปรแกรมตรวจสอบว่าสามารถระบายสีต้นไม้ให้โหนดติดกันไม่ซ้ำสีได้หรือไม่ใน Python

สมมติว่าต้นไม้แบบทวิภาคมีค่าโหนดแทนสี (สูงสุด 2 สี) ตรวจสอบว่าสลับสีได้หรือไม่ เพื่อให้โหนดติดกันไม่มีสีเหมือนกัน

เช่น ต้นไม้ที่มีโหนดสี 2 และ 1 สามารถสลับได้เป็นรูปแบบสลับกัน

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

• colors := นับความถี่แต่ละสี
• prop := นับจำนวนโหนดในระดับคู่/คี่ (flag True/False)
• DFS ผ่านต้นไม้ สลับ flag ทุกระดับ
• ถ้าชุดค่าของ colors และ prop เท่ากัน คืน True

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

from collections import defaultdict
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):
        colors = defaultdict(int)
        prop = defaultdict(int)
        def dfs(node, flag=True):
            if not node: return
            colors[node.val] += 1
            prop[flag] += 1
            dfs(node.left, not flag)
            dfs(node.right, not flag)
        dfs(root)
        return set(colors.values()) == set(prop.values())

ob = Solution()
root = TreeNode(2)
root.left = TreeNode(2)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
root.right.left.right = TreeNode(1)
print(ob.solve(root))

อินพุต

root = TreeNode(2) ...

เอาต์พุต

True