Update bt_lowest_common_ancestor.py

trees/bt_lowest_common_ancestor.py

@@ -3,25 +3,27 @@
 # author: bt3gl
 
 
-class Tree:
+def lowest_common_ancestor(root, p, q):
+
+    stack = [root]
+    parent = {root: None}
     
-    def lowest_common_ancestor(self, root, p, q):
+    while p not in parent or q not in parent:
-        
-        def dfs(root, p, q):
 
-            if not root:
+        node = stack.pop()
-                return False
+        if node:
-            
+            parent[node.left] = node
-            left = dfs(root.left, p, q)
+            parent[node.right] = node
-            right = dfs(root.right, p, q)
+            stack.append(node.left)
-            mid = root == p or root == q
+            stack.append(node.right)
 
-            if mid + left + right >= 2:
+    ancestors = set()
-                self.answer = root
+    while p:
+        ancestors.add(p)
+        p = parent[p]
+
+    while q not in ancestors:
+        q = parent[q]
             
-            return left or right or mid
+    return q
                 
-        dfs(root, p, q)
-        
-        return self.answer
-