Merge pull request #741 from openset/develop

Add: Balanced Binary Tree

Author: Shuo

problems/balanced-binary-tree/balanced_binary_tree.go

@@ -14,26 +14,25 @@ type TreeNode = kit.TreeNode
  * }
  */
 func isBalanced(root *TreeNode) bool {
-	_, isBalanced := recur(root)
-	return isBalanced
-}
-
-func recur(root *TreeNode) (int, bool) {
-	if root == nil {
-		return 0, true
-	}
-	leftDepth, leftIsBalanced := recur(root.Left)
-	rightDepth, rightIsBalanced := recur(root.Right)
-	if leftIsBalanced && rightIsBalanced &&
-		-1 <= leftDepth-rightDepth && leftDepth-rightDepth <= 1 {
-		return max(leftDepth, rightDepth) + 1, true
+	if root != nil {
+		left := depth(root.Left)
+		right := depth(root.Right)
+		if left > right+1 || left < right-1 {
+			return false
+		}
+		return isBalanced(root.Left) && isBalanced(root.Right)
 	}
-	return 0, false
+	return true
 }
 
-func max(a, b int) int {
-	if a > b {
-		return a
+func depth(root *TreeNode) int {
+	if root != nil {
+		left := depth(root.Left) + 1
+		right := depth(root.Right) + 1
+		if left > right {
+			return left
+		}
+		return right
 	}
-	return b
+	return 0
 }

problems/balanced-binary-tree/balanced_binary_tree_test.go

@@ -1 +1,31 @@
 package problem110
+
+import (
+	"testing"
+
+	"github.com/openset/leetcode/internal/kit"
+)
+
+type testType struct {
+	in   []int
+	want bool
+}
+
+func TestIsBalanced(t *testing.T) {
+	tests := [...]testType{
+		{
+			in:   []int{3, 9, 20, kit.NULL, kit.NULL, 15, 7},
+			want: true,
+		},
+		{
+			in:   []int{1, 2, 2, 3, 3, kit.NULL, kit.NULL, 4, 4},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		got := isBalanced(kit.SliceInt2TreeNode(tt.in))
+		if got != tt.want {
+			t.Fatalf("in: %v, got: %v, want: %v", tt.in, got, tt.want)
+		}
+	}
+}