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

Suppose we have a binary tree and another value k, we have to find the number of unique node to sub child paths are there which sums to k.

So, if the input is like

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

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() . This will take node
• if node is not null, then
  • prefix := prefix + value of node
  • ans := ans + (count[prefix - target], if this is not available, it will be 0)
  • 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, do the following
• dfs(root)
• return ans

Let us see the following implementation to get better understanding −

Example 

Live Demo

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

Input

root = TreeNode(3)
root.left = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(1)
root.right.left.right = TreeNode(2)
5

Output

2

Arnab Chakraborty

Published on 02-Dec-2020 05:15:50

• Related Questions & Answers
• Program to count number of fraction pairs whose sum is 1 in python
• Program to find number of consecutive subsequences whose sum is divisible by k in Python
• Program to count number of consecutive lists whose sum is n in C++
• Program to count number of paths with cost k from start to end point in Python
• Program to find sum of rectangle whose sum at most k in Python
• Program to find number of sublists whose sum is given target in python
• Program to find sum of k non-overlapping sublists whose sum is maximum in C++
• Count pairs in array whose sum is divisible by K in C++
• C++ Program to find Number Whose XOR Sum with Given Array is a Given Number k
• Find the Minimum Number of Fibonacci Numbers Whose Sum Is K in C++
• Program to count number of unique paths that includes given edges in Python
• Program to find three unique elements from list whose sum is closest to k Python
• Program to find max number of K-sum pairs in Python
• Find number of substrings of length k whose sum of ASCII value of characters is divisible by k in C++
• Program to count subsets that sum up to k in python