diff --git a/README.md b/README.md index 170fce6..480dcba 100644 --- a/README.md +++ b/README.md @@ -130,3 +130,4 @@ problems from * [Day 120](https://github.com/vaskoz/dailycodingproblem-go/issues/248) * [Day 123](https://github.com/vaskoz/dailycodingproblem-go/issues/253) * [Day 124](https://github.com/vaskoz/dailycodingproblem-go/issues/255) +* [Day 125](https://github.com/vaskoz/dailycodingproblem-go/issues/257) diff --git a/day125/problem.go b/day125/problem.go new file mode 100644 index 0000000..dba6a4b --- /dev/null +++ b/day125/problem.go @@ -0,0 +1,45 @@ +package day125 + +// BST is an integer binary search tree. +type BST struct { + Val int + Left, Right *BST +} + +// FindTwoSumKInBST returns a slice of integer nodes that sum to K. +// If no two nodes exist that sum to K, then return nil. +// Runs in O(N log N) where N is the size of the BST. +func FindTwoSumKInBST(head *BST, k int) []int { + return inorder(head, head, k, []int{}) +} + +func inorder(head, current *BST, k int, result []int) []int { + if current == nil { + return result + } + leftResult := inorder(head, current.Left, k, result) + if len(leftResult) == 2 { + return leftResult + } + node := searchBST(head, k-current.Val) + if node != nil && node != current { + return []int{current.Val, node.Val} + } + rightResult := inorder(head, current.Right, k, result) + if len(rightResult) == 2 { + return rightResult + } + return result +} + +func searchBST(head *BST, target int) *BST { + if head == nil { + return nil + } + if target == head.Val { + return head + // } else if target < head.Val { // this part of searchBST is unused due to the + // return searchBST(head.Left, target) // inorder traversal that calls this search. + } + return searchBST(head.Right, target) +} diff --git a/day125/problem_test.go b/day125/problem_test.go new file mode 100644 index 0000000..2952ea6 --- /dev/null +++ b/day125/problem_test.go @@ -0,0 +1,42 @@ +package day125 + +import ( + "reflect" + "testing" +) + +var testcases = []struct { + head *BST + k int + expected []int +}{ + {&BST{10, &BST{5, nil, nil}, &BST{ + 15, &BST{11, nil, nil}, &BST{15, nil, nil}, + }}, 20, []int{5, 15}}, + {&BST{10, &BST{5, nil, nil}, &BST{ + 15, &BST{11, nil, nil}, &BST{15, nil, nil}, + }}, 100, []int{}}, + {&BST{10, &BST{5, nil, nil}, &BST{ + 15, &BST{11, nil, nil}, &BST{15, nil, nil}, + }}, 15, []int{5, 10}}, + {&BST{10, &BST{5, nil, nil}, &BST{ + 15, &BST{11, nil, nil}, &BST{15, nil, nil}, + }}, 26, []int{11, 15}}, +} + +func TestFindTwoSumKInBST(t *testing.T) { + t.Parallel() + for _, tc := range testcases { + if result := FindTwoSumKInBST(tc.head, tc.k); !reflect.DeepEqual(result, tc.expected) { + t.Errorf("Expected %v got %v", tc.expected, result) + } + } +} + +func BenchmarkFindTwoSumInBST(b *testing.B) { + for i := 0; i < b.N; i++ { + for _, tc := range testcases { + FindTwoSumKInBST(tc.head, tc.k) + } + } +}