Golang O(n) solution using map with some explanation


  • 0

    I think the most important thing we should notice to come up with this optimal solution is,
    We can calculate any sum of partial paths from A to B by (sum from root to B) - (sum from root to A).
    So just by storing all sums from root to each node, we can calculate any sum of possible path in O(1).

    If we understand that, only thing we need to do is just traverse a node from root, store all sums from root to map, then check if there is any pair found like Two Sum problem.

    func pathSum(root *TreeNode, sum int) int {
    	if root == nil {
    		return 0
    	}
    
    	sumMap := make(map[int]int)
    	sumMap[0] = 1
    	num := 0
    	doPathSum(&num, root, 0, sum, sumMap)
    	return num
    }
    
    func doPathSum(res *int, node *TreeNode, sumSoFar int, target int, cache map[int]int) {
    	if node == nil {
    		return
    	}
    
    	sumSoFar += node.Val
    	if cnt, ok := cache[sumSoFar-target]; ok {
    		*res += cnt
    	}
    
    	if cnt, ok := cache[sumSoFar]; ok {
    		cache[sumSoFar] = cnt + 1
    	} else {
    		cache[sumSoFar] = 1
    	}
    
    	doPathSum(res, node.Left, sumSoFar, target, cache)
    	doPathSum(res, node.Right, sumSoFar, target, cache)
    
    	cache[sumSoFar]--
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.