diff --git a/problems/cousins-in-binary-tree/cousins_in_binary_tree.go b/problems/cousins-in-binary-tree/cousins_in_binary_tree.go new file mode 100644 index 000000000..3dd17d295 --- /dev/null +++ b/problems/cousins-in-binary-tree/cousins_in_binary_tree.go @@ -0,0 +1,40 @@ +package cousins_in_binary_tree + +import . "github.com/openset/leetcode/internal/kit" + +/** + * Definition for a binary tree node. + * type TreeNode struct { + * Val int + * Left *TreeNode + * Right *TreeNode + * } + */ +func isCousins(root *TreeNode, x int, y int) bool { + xp, yp, xd, yd, depth := 0, 0, 0, 0, 0 + queue := []*TreeNode{root} + for l := len(queue); l > 0 && xd == 0 && yd == 0; l = len(queue) { + for _, node := range queue { + if node.Left != nil { + if node.Left.Val == x { + xp, xd = node.Val, depth+1 + } else if node.Left.Val == y { + yp, yd = node.Val, depth+1 + } else { + queue = append(queue, node.Left) + } + } + if node.Right != nil { + if node.Right.Val == x { + xp, xd = node.Val, depth+1 + } else if node.Right.Val == y { + yp, yd = node.Val, depth+1 + } else { + queue = append(queue, node.Right) + } + } + } + depth, queue = depth+1, queue[l:] + } + return xp != yp && xd == yd +} diff --git a/problems/cousins-in-binary-tree/cousins_in_binary_tree_test.go b/problems/cousins-in-binary-tree/cousins_in_binary_tree_test.go new file mode 100644 index 000000000..798711218 --- /dev/null +++ b/problems/cousins-in-binary-tree/cousins_in_binary_tree_test.go @@ -0,0 +1,55 @@ +package cousins_in_binary_tree + +import ( + "testing" + + . "github.com/openset/leetcode/internal/kit" +) + +type caseType struct { + input []int + x int + y int + expected bool +} + +func TestIsCousins(t *testing.T) { + tests := [...]caseType{ + { + input: []int{1, 2, 3, 4}, + x: 4, + y: 3, + expected: false, + }, + { + input: []int{1, 2, 3, NULL, 4, NULL, 5}, + x: 5, + y: 4, + expected: true, + }, + { + input: []int{1, 2, 3, NULL, 4}, + x: 2, + y: 3, + expected: false, + }, + { + input: []int{1, 2, 3, NULL, 4, 5}, + x: 4, + y: 5, + expected: true, + }, + { + input: []int{1, 2, 3, NULL, 4, 5}, + x: 5, + y: 4, + expected: true, + }, + } + for _, tc := range tests { + output := isCousins(SliceInt2TreeNode(tc.input), tc.x, tc.y) + if output != tc.expected { + t.Fatalf("input: %v, output: %v, expected: %v", tc.input, output, tc.expected) + } + } +}