educative.io

A proposal for improving the time complexity

It seems that instead of looping over the the path at each node, a sliding window can be used. So the worst case becomes O(nlogn)

Proposed solution:

def count_paths(root, S):
  if root is None:
    return 0
  return count_paths_rec(root, S, [], 0, 0)

def count_paths_rec(node, S, cur_path, cur_sum, start):
  cur_path.append(node.val)
  cur_sum += node.val

  # advance start of path
  while cur_sum > S and start < len(cur_path)-1:
    cur_sum -= cur_path[start]
    start += 1

  count = 0
  if cur_sum == S:
    count += 1
  
  for child in (node.left, node.right):
    if child is not None:
      count += count_paths_rec(child, S, cur_path, cur_sum, start)

  cur_path.pop()

  return count

Please let me know if there’s something I’m missing :slight_smile:

1 Like

I think the whole path sum might be wrong when we move the start in case we meet negative numbers or zero in the next recursion, ex : S = 3, path = 3 0 1 -1

Yes, you’re quite right!
In this case a different optimization can be made: for each path a hash table would be maintained containing the prefix sums for that path (counting the number of occurences for each possible prefix sum). Then, whenever a new node in the path is reached, we know the sum of the entire path up to this node, and can check the hash table for prefixes such that if we remove them we get the desired sum.
The hash table would be updating via backtracking.