algorithms/leetcode/437-PathSumIII.go at master · freewu/algorithms · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package main

// 437. Path Sum III
// Given the root of a binary tree and an integer targetSum,
// return the number of paths where the sum of the values along the path equals targetSum.

// The path does not need to start or end at the root or a leaf,
// but it must go downwards (i.e., traveling only from parent nodes to child nodes).

// Example 1:
// <img src="https://assets.leetcode.com/uploads/2021/04/09/pathsum3-1-tree.jpg" / >
// Input: root = [10,5,-3,3,2,null,11,3,-2,null,1], targetSum = 8
// Output: 3
// Explanation: The paths that sum to 8 are shown.

// Example 2:
// Input: root = [5,4,8,11,null,13,4,7,2,null,null,5,1], targetSum = 22
// Output: 3

// Constraints:
//     The number of nodes in the tree is in the range [0, 1000].
//     -10^9 <= Node.val <= 10^9
//     -1000 <= targetSum <= 1000

import "fmt"

// Definition for a binary tree node.
type TreeNode struct {
    Val int
    Left *TreeNode
    Right *TreeNode
}

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
// dfs
func pathSum(root *TreeNode, targetSum int) int {
    res := 0
    var dfs func(tree *TreeNode, targetSum int, pathStack []*TreeNode)
    dfs = func (tree *TreeNode, targetSum int, pathStack []*TreeNode) {
        if tree == nil {
            return
        }
        sum := 0
        pathStack = append(pathStack, tree)
        for i := len(pathStack) - 1; i >= 0; i-- {
            sum += pathStack[i].Val
            if sum == targetSum {
                res++
            }
        }
        dfs(tree.Left, targetSum, pathStack)
        dfs(tree.Right, targetSum, pathStack)
    }
    dfs(root, targetSum, []*TreeNode{})
    return res
}

func pathSum1(root *TreeNode, targetSum int) int {
    res, m := 0, map[int]int{0:1}
    var dfs func(*TreeNode, int)
    dfs = func(root *TreeNode, cur int) {
        if root == nil {
            return
        }
        cur += root.Val
        res += m[cur - targetSum]
        m[cur]++
        defer func() {m[cur]--}()
        dfs(root.Left, cur)
        dfs(root.Right, cur)
    }
    dfs(root, 0)
    return res
}

func main() {
    tree1 := &TreeNode {
        10,
        &TreeNode {
            5,
            &TreeNode {
                3,
                &TreeNode{3, nil, nil},
                &TreeNode{-2, nil, nil},
            },
            &TreeNode {
                2,
                nil,
                &TreeNode{1, nil, nil},
            },
        },
        &TreeNode {
            -3,
            nil,
            &TreeNode{11, nil, nil},
        },
    }
    tree2 := &TreeNode {
        5,
        &TreeNode {
            4,
            &TreeNode {
                11,
                &TreeNode{7, nil, nil},
                &TreeNode{2, nil, nil},
            },
            nil,
        },
        &TreeNode {
            8,
            &TreeNode{13, nil, nil},
            &TreeNode{
                4,
                &TreeNode{5, nil, nil},
                &TreeNode{1, nil, nil},
            },
        },
    }
    fmt.Println(pathSum(tree1,8)) // 3
    fmt.Println(pathSum(tree2,22)) // 3

    fmt.Println(pathSum1(tree1,8)) // 3
    fmt.Println(pathSum1(tree2,22)) // 3
}