67 lines
1.3 KiB
Go
67 lines
1.3 KiB
Go
package top100liked
|
|
|
|
import (
|
|
"fmt"
|
|
"testing"
|
|
)
|
|
|
|
// https://leetcode.cn/problems/path-sum-iii/description/?envType=study-plan-v2&envId=top-100-liked
|
|
|
|
type TreeNode struct {
|
|
Val int
|
|
Left *TreeNode
|
|
Right *TreeNode
|
|
}
|
|
|
|
// 朴素解法，可以AC，但时间复杂度高
|
|
// func pathSum(root *TreeNode, targetSum int) int {
|
|
// if root == nil {
|
|
// return 0
|
|
// }
|
|
// res := countFrom(root, targetSum)
|
|
// res += pathSum(root.Left, targetSum)
|
|
// res += pathSum(root.Right, targetSum)
|
|
// return res
|
|
// }
|
|
// func countFrom(node *TreeNode, target int) int {
|
|
// if node == nil {
|
|
// return 0
|
|
// }
|
|
// c := 0
|
|
// if node.Val == target {
|
|
// c++
|
|
// }
|
|
// c += countFrom(node.Left, target-node.Val)
|
|
// c += countFrom(node.Right, target-node.Val)
|
|
// return c
|
|
// }
|
|
|
|
// 前缀和
|
|
func pathSum(root *TreeNode, targetSum int) int {
|
|
m := make(map[int]int)
|
|
count := 0
|
|
m[0] = 1
|
|
var dfs func(*TreeNode, int)
|
|
dfs = func(root *TreeNode, cur int) {
|
|
if root == nil {
|
|
return
|
|
}
|
|
cur += root.Val
|
|
count += m[cur-targetSum]
|
|
m[cur]++
|
|
dfs(root.Left, cur)
|
|
dfs(root.Right, cur)
|
|
m[cur]--
|
|
}
|
|
dfs(root, 0)
|
|
return count
|
|
}
|
|
|
|
func Test(t *testing.T) {
|
|
input := &TreeNode{
|
|
Val: 1,
|
|
Right: &TreeNode{Val: 2, Right: &TreeNode{Val: 3, Right: &TreeNode{Val: 4, Right: &TreeNode{Val: 5}}}},
|
|
}
|
|
fmt.Println(pathSum(input, 3) == 2)
|
|
}
|