Program to count number of paths whose sum is k in python

In this problem, we need to find the number of paths in a binary tree that sum to a target value k. The paths can start and end at any nodes, not necessarily from root to leaf.

So, if the input is like:

and k = 5, then the output will be 2, as the paths are [2, 3] and [1, 4].

Algorithm

To solve this, we will follow these steps ?

• count := a map initially holds value 1 for key 0
• ans := 0, prefix := 0
• Define a function dfs() that takes a node
• if node is not null, then:
  • prefix := prefix + value of node
  • ans := ans + count[prefix − target]
  • count[prefix] := count[prefix] + 1
  • dfs(left of node)
  • dfs(right of node)
  • count[prefix] := count[prefix] − 1
  • prefix := prefix − value of node
• From the main method: dfs(root) and return ans

How It Works

The algorithm uses prefix sum and backtracking. We maintain a running sum (prefix) and count how many times each prefix sum has occurred. For each node, we check if there's a previous prefix sum that, when subtracted from current prefix, equals our target.

Example

Let us see the following implementation to get better understanding ?

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)
                
                # Backtrack
                count[prefix] -= 1
                prefix -= node.val
                
        dfs(root)
        return ans

# Create the tree
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
result = ob.solve(root, k)
print(f"Number of paths with sum {k}: {result}")

The output of the above code is ?

Number of paths with sum 5: 2

Key Points

• Prefix Sum: We maintain cumulative sum from root to current node
• Counter: Tracks frequency of each prefix sum encountered
• Backtracking: We restore the state after exploring each subtree
• Path Detection: If prefix − target exists in counter, we found valid paths

Conclusion

This solution uses prefix sum with backtracking to efficiently count paths summing to target k. The time complexity is O(n) where n is the number of nodes, making it optimal for this problem.