#!/usr/bin/env python
# coding: utf-8

# <h1>236. Lowest Common Ancestor of a Binary Tree</h1>
# <hr>
# 
# <!--Copy Paste Leetcode statement between-->
# <p>Given a binary tree, find the lowest common ancestor (LCA) of two given nodes in the tree.</p>
# 
# <p>According to the <a href="https://en.wikipedia.org/wiki/Lowest_common_ancestor" target="_blank">definition of LCA on Wikipedia</a>: “The lowest common ancestor is defined between two nodes p&nbsp;and q&nbsp;as the lowest node in T that has both p&nbsp;and q&nbsp;as descendants (where we allow <b>a node to be a descendant of itself</b>).”</p>
# 
# <p>&nbsp;</p>
# <p><strong>Example 1:</strong></p>
# <img alt="" src="./img1.png" style="width: 200px; height: 190px;">
# <pre><strong>Input:</strong> root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 1
# <strong>Output:</strong> 3
# <strong>Explanation:</strong> The LCA of nodes 5 and 1 is 3.
# </pre>
# 
# <p>&nbsp;</p>
# <p><strong>Example 2:</strong></p>
# <img alt="" src="./img2.png" style="width: 200px; height: 190px;">
# <pre><strong>Input:</strong> root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 4
# <strong>Output:</strong> 5
# <strong>Explanation:</strong> The LCA of nodes 5 and 4 is 5, since a node can be a descendant of itself according to the LCA definition.
# </pre>
# 
# <p>&nbsp;</p>
# <p><strong>Example 3:</strong></p>
# 
# <pre><strong>Input:</strong> root = [1,2], p = 1, q = 2
# <strong>Output:</strong> 1
# </pre>
# 
# <p>&nbsp;</p>
# <p><strong>Constraints:</strong></p>
# 
# <ul>
# 	<li>The number of nodes in the tree is in the range <code>[2, 10<sup>5</sup>]</code>.</li>
# 	<li><code>-10<sup>9</sup> &lt;= Node.val &lt;= 10<sup>9</sup></code></li>
# 	<li>All <code>Node.val</code> are <strong>unique</strong>.</li>
# 	<li><code>p != q</code></li>
# 	<li><code>p</code> and <code>q</code> will&nbsp;exist in the tree.</li>
# </ul>
# 
# <!--Copy Paste Leetcode statement between-->
# 
# <p>&nbsp;</p>
# <a href="https://leetcode.com/problems/lowest-common-ancestor-of-a-binary-tree/">Source</a> 
# <hr>

# In[1]:


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


# <h4>Code</h4>

# In[ ]:


def lowest_common_ancestor(root, p, q):
    """Recursive bottom-up solution using a hashmap of parent pointers"""
    def dfs(node, parent):
        if not node:
            return
        if p in parent and q in parent:
            return
        if node.left:
            parent[node.left] = node
            dfs(node.left, parent)
        if node.right:
            parent[node.right] = node
            dfs(node.right, parent)

    parent = {root:None}
    dfs(root, parent)
    ancestors = set()
    while p:                  # find all ancestors of p
        ancestors.add(p)
        p = parent[p]
    while q not in ancestors: # while nothing common
        q = parent[q]
    return q                  # 1st one to be also an ancestors of p


# In[ ]:


def lowest_common_ancestor(root, p, q):
    """Alternative recursive bottom-up solution"""
    def dfs(node):
        nonlocal lca
        if not node:
            return False         # return None is not acceptable because addition below
        left = dfs(node.left)
        right = dfs(node.right)

        curr = True if node in (p, q) else False
        if curr + left + right >= 2:     # If any two variables is True
            lca = node
        return curr or left or right     # return the one where we found something

    lca = None
    dfs(root)
    return lca


# In[ ]:


def lowest_common_ancestor(root, p, q):
    """Another alternative recursive bottom-up solution"""
    if not root:
        return False             # return None is acceptable here
    if root in (p, q):
        return root

    # Option A - a bit faster
    left = right = None
    if root.left:
        left = lowest_common_ancestor(root, p, q)
    if root.right:
        right = lowest_common_ancestor(root, p, q)

    # Option B - a bit slower
    # left = lowest_common_ancestor(root, p, q)
    # right = lowest_common_ancestor(root, p, q)

    if left and right:    # if we found something on each side
        return root
    return left or right  # otherwise, return the one where we found something


# <hr>
# <h4>Follow up:</h4>
# <p>Solve it both recursively and iteratively.</p>
# 
# <h4>Code</h4>

# In[ ]:


def lowest_common_ancestor(root, p, q):
    """Iterative version of the first solution (with hashmap of parent pointers)"""
    stack = [root]
    parent = {root: None}
    while not (p in parent and q in parent):
        node = stack.pop()
        if node.left:
            parent[node.left] = node
            stack.append(node.left)
        if node.right:
            parent[node.right] = node
            stack.append(node.right)
    ancestors = set()
    while p:
        ancestors.add(p)
        p = parent[p]
    while q not in ancestors:
        q = parent[q]
    return q