update algorithm as per channel

go/0110-balanced-binary-tree.go

@@ -1,26 +1,34 @@
+/**
+ * Definition for a binary tree node.
+ * type TreeNode struct {
+ * Val int
+ * Left *TreeNode
+ * Right *TreeNode
+ * }
+ */
+type BalanceTree struct {
+ Balance bool
+ Height int
+}
+
 func isBalanced(root *TreeNode) bool {
- return isBalancedUtil(root) != math.MaxUint32
+ return dfs(root).Balance
 }
 
-func isBalancedUtil(root *TreeNode) int {
+func dfs(root *TreeNode) BalanceTree {
  if root == nil {
- return 0
+ return BalanceTree{true, 0}
  }
 
- left := isBalancedUtil(root.Left)
- right := isBalancedUtil(root.Right)
-
- if left == math.MaxUint32 || right == math.MaxUint32 {
- return math.MaxUint32
- }
+ left, right := dfs(root.Left), dfs(root.Right)
+ balanced := (left.Balance && right.Balance && int(math.Abs(float64(left.Height)-float64(right.Height))) <= 1)
 
- if math.Abs(float64(left-right)) > 1 {
- return math.MaxUint32
- }
+ return BalanceTree{balanced, 1 + max(left.Height, right.Height)}
+}
 
- if left > right {
- return left + 1
+func max(a, b int) int {
+ if a > b {
+ return a
  }
-
- return right + 1
+ return b
 }