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