feature: balanced binary tree dfs revisit

trees/BalancedBinaryTree.py

@@ -4,7 +4,7 @@
 
 
 # O(n) time || O(n) space
-def is_balanced(self, root: Optional[TreeNode]) -> bool:
+def is_balanced_rec(self, root: Optional[TreeNode]) -> bool:
     def dfs(node):
         if not node:
             return 0, True
@@ -15,3 +15,29 @@ def dfs(node):
         return max(l_height, r_height) + 1, l_balanced and r_balanced and abs(l_height - r_height) <= 1
 
     return dfs(root)[1]
+
+
+# O(n) time || O(n) space
+def is_balanced_dfs(self, root: Optional[TreeNode]) -> bool:
+    stack = [(root, 0, False)]
+    heights = {}
+    while stack:
+        node, height, visited = stack.pop()
+
+        if not node:
+            continue
+
+        if visited:
+            l_height = heights.get(node.left, 0)
+            r_height = heights.get(node.right, 0)
+
+            if abs(l_height - r_height) > 1:
+                return False
+
+            heights[node] = 1 + max(l_height, r_height)
+        else:
+            stack.append((node, height, True))
+            stack.append((node.left, height + 1, False))
+            stack.append((node.right, height + 1, False))
+
+    return True

trees/test/test_is_balanced.py

@@ -1,6 +1,6 @@
 import unittest
 
-from trees.BalancedBinaryTree import is_balanced
+from trees.BalancedBinaryTree import is_balanced_rec, is_balanced_dfs
 from trees.TreeNode import TreeNode
 
 
@@ -15,7 +15,8 @@ def test_is_balanced(self):
                 right=TreeNode(7)
             )
         )
-        self.assertEqual(True, is_balanced(self, root))
+        self.assertEqual(True, is_balanced_rec(self, root))
+        self.assertEqual(True, is_balanced_dfs(self, root))
 
     def test_is_balanced_1(self):
         root = TreeNode(
@@ -31,7 +32,8 @@ def test_is_balanced_1(self):
             ),
             right=TreeNode(2)
         )
-        self.assertEqual(False, is_balanced(self, root))
+        self.assertEqual(False, is_balanced_rec(self, root))
+        self.assertEqual(False, is_balanced_dfs(self, root))
 
 
 if __name__ == '__main__':