From b51545e39af4a7d47bd83b31f61437d2f699d9e3 Mon Sep 17 00:00:00 2001 From: openset Date: Fri, 22 Feb 2019 10:44:26 +0800 Subject: [PATCH 1/2] Add: Cousins in Binary Tree --- .../cousins_in_binary_tree.go | 40 +++++++++++++++++ .../cousins_in_binary_tree_test.go | 43 +++++++++++++++++++ 2 files changed, 83 insertions(+) create mode 100644 problems/cousins-in-binary-tree/cousins_in_binary_tree.go create mode 100644 problems/cousins-in-binary-tree/cousins_in_binary_tree_test.go 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..511e1c180 --- /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; 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..e2702e33d --- /dev/null +++ b/problems/cousins-in-binary-tree/cousins_in_binary_tree_test.go @@ -0,0 +1,43 @@ +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, + }, + } + 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) + } + } +} From 9a5c65dbd0d729a07cd0b1f5b2ad884deba00a9a Mon Sep 17 00:00:00 2001 From: openset Date: Fri, 22 Feb 2019 11:08:30 +0800 Subject: [PATCH 2/2] Add: test case --- .../cousins-in-binary-tree/cousins_in_binary_tree.go | 2 +- .../cousins_in_binary_tree_test.go | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/problems/cousins-in-binary-tree/cousins_in_binary_tree.go b/problems/cousins-in-binary-tree/cousins_in_binary_tree.go index 511e1c180..3dd17d295 100644 --- a/problems/cousins-in-binary-tree/cousins_in_binary_tree.go +++ b/problems/cousins-in-binary-tree/cousins_in_binary_tree.go @@ -13,7 +13,7 @@ import . "github.com/openset/leetcode/internal/kit" 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; l = len(queue) { + 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 { 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 index e2702e33d..798711218 100644 --- a/problems/cousins-in-binary-tree/cousins_in_binary_tree_test.go +++ b/problems/cousins-in-binary-tree/cousins_in_binary_tree_test.go @@ -33,6 +33,18 @@ func TestIsCousins(t *testing.T) { 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)