Merge pull request #382 from openset/develop

Add: Cousins in Binary Tree

Author: Shuo

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

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