โปรแกรมนับจำนวนเส้นทางที่ผลรวมเท่ากับ k ในต้นไม้แบบทวิภาคใน Python

สมมติว่าเรามีต้นไม้แบบทวิภาคและค่า k หาจำนวนเส้นทางจากโหนดไปยังโหนดลูกที่ผลรวมเท่ากับ k

เช่น ต้นไม้ที่มีโหนด 3,2,4,1,2 และ k=5 ให้ผลลัพธ์ 2 เส้นทาง [2,3] และ [1,4]

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

• ใช้ Counter เก็บความถี่ของ prefix sum เริ่มต้นด้วย {0:1}
• ans := 0, prefix := 0
• DFS ผ่านต้นไม้:
  • prefix += node.val
  • ans += count[prefix - target]
  • count[prefix] += 1
  • DFS ซ้ายและขวา
  • Backtrack: count[prefix] -= 1, prefix -= node.val
• คืนค่า ans

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

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

อินพุต

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

เอาต์พุต

2