Iterative Solutions in Python/C++


  • 32

    Solution 1

    Same algorithm as my recursive solution (look there if you want some explanation), but iterative. I do a post-order traversal with a stack. Each stack element at first is a [node, parent] pair, where parent is the stack element of the node's parent node. When the children of a parent get finished, their results are appended to their parent's stack element. So when a parent gets finished, we have the results of its children/subtrees available (its stack element at that point is [node, parent, resultForLeftSubtree, resultForRightSubtree]).

    def lowestCommonAncestor(self, root, p, q):
        answer = []
        stack = [[root, answer]]
        while stack:
            top = stack.pop()
            (node, parent), subs = top[:2], top[2:]
            if node in (None, p, q):
                parent += node,
            elif not subs:
                stack += top, [node.right, top], [node.left, top]
            else:
                parent += node if all(subs) else max(subs),
        return answer[0]
    

    Solution 2

    Here I find the paths to p and q and then find the last node where the paths match. I just came up with the path-building technique for this, and I find it quite neat and maybe it could be useful elsewhere.

    def lowestCommonAncestor(self, root, p, q):
        def path(root, goal):
            path, stack = [], [root]
            while True:
                node = stack.pop()
                if node:
                    if node not in path[-1:]:
                        path += node,
                        if node == goal:
                            return path
                        stack += node, node.right, node.left
                    else:
                        path.pop()
        return next(a for a, b in zip(path(root, p), path(root, q))[::-1] if a == b)
    

    C++ version of Solution 1

    I don't use C++ much, so maybe there's room for improvement with stuff that I don't know.

    class Solution {
        struct Frame {
            TreeNode* node;
            Frame* parent;
            vector<TreeNode*> subs;
        };
    public:
        TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
            Frame answer;
            stack<Frame> stack;
            stack.push({root, &answer});
            while (stack.size()) {
                Frame *top = &stack.top(), *parent = top->parent;
                TreeNode *node = top->node;
                if (!node || node == p || node == q) {
                    parent->subs.push_back(node);
                    stack.pop();
                } else if (top->subs.empty()) {
                    stack.push({node->right, top});
                    stack.push({node->left, top});
                } else {
                    TreeNode *left = top->subs[0], *right = top->subs[1];
                    parent->subs.push_back(!left ? right : !right ? left : node);
                    stack.pop();
                }
            }
            return answer.subs[0];
        }
    };

  • 0
    G

    Could you please rewrite the Solution 1 in C++. Many thanks!


  • 0

    Hi, greendreams. I guess the first solution is very Pythonic (especially the nested lists, the reference and the commas ,) and may be hard to be translated into C++.

    BTW, the second path-finding solution is really unique. It seems that few people have this idea.


  • 0

    Ok, I added a C++ version. I don't use C++ much, so maybe there's room for improvement with stuff that I don't know.


  • 0
    Z

    One minor comment: to make sure that top->subs[0] is the result of the left subtree and top->subs[1] is the result of the right subtree, we should first push node->right and then node->left within the following "else if" statement.

           else if (top->subs.empty()) {
                stack.push({node->left, top});
                stack.push({node->right, top});
            }
    

    However, since left and right are symmetric regarding this problem, the current solution will give us the same result.


  • 0

    Thanks a lot! I fixed it now.


  • 0

    Wow, originally I think it would be hard to turn the first solution into C++. You just make it by defining the nice struct Frame :-)


  • 0

    Nothing special, though, pretty much just a stack frame like in normal recursion, but done myself. After all, I am just simulating recursive calls there.


  • 0
    D

    I must say that your code is very neat and hard to invent. :) I came up the same recursive solution as yours but had a hard time trying to translate it into iterative one. The main difficulty is to pass the result up.

    I finally succeeded. The main idea is to use two variables to track the ancestors of p and q. If at a point we find a node's children consist of both ancestors, that node is the LCA. If its children consist of only one ancestor, update the ancestor to the node itself. It's much longer with a running time of ~200 ms. Any suggestion to make it shorter while keeping (almost) the same speed, or faster?

    class Solution:
        # @param {TreeNode} root
        # @param {TreeNode} p
        # @param {TreeNode} q
        # @return {TreeNode}
        def lowestCommonAncestor(self, root, p, q):
            stack = [(root, False)]
            ancestor_p, ancestor_q, pFound, qFound = root, root, False, False
            while stack:
                node, visited = stack.pop()
                if node:
                    if node is p:
                        ancestor_p, pFound = node, True
                    elif node is q:
                        ancestor_q, qFound = node, True
                    elif not visited:
                        # post-order traversal
                        stack.append((node, True))
                        stack.append((node.right, False))
                        stack.append((node.left, False))
                    else:
                        # left subtree and right subtree visited
                        if (node.left is ancestor_p and node.right is ancestor_q) or (node.left is ancestor_q and node.right is ancestor_p):
                            # node is the junction of ancestors of p and q
                            return node
                        if (node.left is ancestor_p) or (node.right is ancestor_p):
                            # update ancestor of p 
                            ancestor_p = node
                        elif (node.left is ancestor_q) or (node.right is ancestor_q):
                            # update ancestor of q
                            ancestor_q = node
    
            '''
            If both p and q's ancestors are found, we should have returned within the while loop
            What we've got here is failure to communicate. That is, p and q is under subtree of the other 
            ''' 
            return p if pFound else q
    

  • 7
    X

    my c++ version of Solution 2

    class Solution {
    public:
        TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
            //iterative, path comparing
            if(!root || root==p || root==q) return root;
            vector<TreeNode*> pathp, pathq, temp;
            temp.push_back(root);
            TreeNode* prev=NULL;
            while(pathp.empty() || pathq.empty()){
                root=temp.back();
                if(root==p) pathp=temp;
                if(root==q) pathq=temp;
                if(!root->left && !root->right || !root->right && prev==root->left || root->right && prev==root->right){
                    temp.pop_back();
                    prev=root;
                }
                else{
                    if(!root->left || prev==root->left) temp.push_back(root->right);
                    else temp.push_back(root->left);
                }
            }
            int n=min(pathp.size(),pathq.size());
             for(int i=1; i<n; i++){
                if(pathp[i]!=pathq[i]) return pathp[i-1];
            }
            return pathp[n-1];
        }
    };

  • 1
    H

    Great Solution! I rewrite your cpp solution to Java version and add comments.
    Your program exactly follow the logic of recursive solution. Good Work!

    public class MyNode{
        TreeNode node;
        MyNode parent;
        boolean visited;
        List<TreeNode> result = new ArrayList<TreeNode>();
        
        public MyNode(TreeNode node, MyNode parent){
            this.node = node;
            this.parent = parent;
        }
    }
    public TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
        MyNode dummy = new MyNode(null, null);//used to retrive global result
        MyNode rootNode = new MyNode(root, dummy);
        Stack<MyNode> stack = new Stack<MyNode>();
        stack.push(rootNode);
        
        while(!stack.isEmpty()){
            MyNode curr = stack.peek();
            TreeNode node = curr.node;
            MyNode parent = curr.parent;
            //if we reach bottom or we found target
            if(node == null || node == p || node == q){
                parent.result.add(node);
                stack.pop();//we are done with this node, pop out
            }else if(!curr.visited){
                //have not visited curr node, push right first then left
                curr.visited = true;
                stack.push(new MyNode(node.right, curr));
                stack.push(new MyNode(node.left, curr));
            }else if(curr.visited){
                //if visited, update result
                TreeNode left = curr.result.get(0);
                TreeNode right = curr.result.get(1);
                if(left != null && right != null){
                    parent.result.add(node);//curr treeNode is LCA
                }else if(left != null){
                    parent.result.add(left);
                }else{
                    parent.result.add(right);
                }
                
                stack.pop();//we are done with this node, pop out
            }
        }
        
        return dummy.result.get(0);
        
    }

  • 0
    G

    @StefanPochmann can you help me out?

    I tried your solution 2 but I wrote your get_path part in a recursive way, and now I get Memory Limit Exceeded for the tree with 20000 nodes. do you have some ways to solve this problem?

    class Solution(object):
        def lowestCommonAncestor(self, root, p, q):
            """
            :type root: TreeNode
            :type p: TreeNode
            :type q: TreeNode
            :rtype: TreeNode
            """
            path1 = []
            self.get_path(root,p,[root],path1)
            path2=[]
            self.get_path(root,q,[root],path2)
            ret = root
            for i in range(0,min(len(path1),len(path2))):
                if path1[i]==path2[i]:
                    ret = path1[i]
            return ret
            
        def get_path(self,root,p,path,res):
            if root==None or len(res):
                return
            if root == p:
                res+=path
                return
            if root.left!=None:
                self.get_path(root.left,p,path+[root.left],res)
                if len(res)!=0:
                    return 
            if root.right!=None:
                self.get_path(root.right,p,path+[root.right],res)

  • 0

    I think the problem is that you create a new path list in every call, and they pile up. The total amount of memory for that is quadratic in the recursion depth, and that tree is very deep. Try extending the path list instead of creating new ones.


  • 0
    G

    oh, you are insightful. now I learned.

    I got AC if I use the reference and change its value before/after each call.

    `

            path.append(root.left)
    
            #self.get_path(root.left,p,path+[root.left],res)
    
            self.get_path(root.left,p,path,res)
    
            path.pop()
    

    `

    thank you ! :)


  • 0
    H

    In Sol1, does the code stop once it reaches one desired node and ignore the subtree with that node as root? I have the doubt since I find that the "if" branch append node to the parent and then do nothing about the children of the node. It will continue on the other branch anyway, but just assume that nothing should be done for the subtree.
    Thanks


  • 0
    N

    C++ Solution

     TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
           map<TreeNode*, TreeNode*> parent;
          stack<TreeNode*> stk;
          parent.insert(pair<TreeNode*,TreeNode*>(root, NULL));
          stk.push(root);
    
          while (parent.find(p)==parent.end() || parent.find(q)==parent.end()) {
              TreeNode *node = stk.top();
              stk.pop();
              if (node->left != NULL) {
                  parent.insert(pair<TreeNode*,TreeNode*>(node->left, node));
                  stk.push(node->left);
              }
              if (node->right != NULL) {
                  parent.insert(pair<TreeNode*,TreeNode*>(node->right, node));
                  stk.push(node->right);
              }
          }
          set<TreeNode*> ancestors;
          while (p != NULL) {
              ancestors.insert(p);
              p = parent.find(p)!=parent.end()? (*(parent.find(p))).second: NULL;
          }
          while (ancestors.find(q)==ancestors.end())
              q = parent.find(q)!=parent.end()? (*(parent.find(q))).second: NULL;
          return q;
        }
    

  • 0
    D

    Here is an iterative solution, which basically mimics the stack operations of the recursive solution:

        struct Vars {
            TreeNode *root, *left, *right;
            int pc;
            Vars(TreeNode *r, int _pc): root(r), left(nullptr), right(nullptr), pc(_pc) {}
        };
    
        TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
            stack<Vars> stk;
            stack<TreeNode*> returns;
            stk.push(Vars(root, 1));
            returns.push(nullptr);
    
            while (!stk.empty()) {
                Vars &v = stk.top();
                TreeNode *tmp;
    
                switch (v.pc) {
                    case 1:
                        if (!v.root || v.root == p || v.root == q) {
                            stk.pop();
                            returns.top() = v.root;
                            break;
                        }
                        v.pc++;
                        break;
                    case 2:
                        returns.push(nullptr);
                        v.pc++;
                        tmp = v.root->left;
                        stk.push(Vars(tmp, 1));
                        break;
                    case 3:
                        v.left = returns.top();
                        v.pc++;
                        tmp = v.root->right;
                        stk.push(Vars(tmp, 1));
                        break;
                    case 4:
                        v.right = returns.top(); returns.pop();
                        if (v.left && v.right) {
                            returns.top() = v.root;
                        } else if (v.left) {
                            returns.top() = v.left;
                        } else {
                            returns.top() = v.right;
                        }
                        stk.pop();
                }
            } // while
            return returns.top();
        }
            
    
        TreeNode* lowestCommonAncestorRecursive(TreeNode* root, TreeNode* p, TreeNode* q) {
            if (!root || root == p || root == q) return root; // Line 1
            auto left = lowestCommonAncestor(root->left, p, q); // Line 2
            auto right = lowestCommonAncestor(root->right, p, q); // Line 3
            if (left && right) return root; // Line 4
            return left ? left : right; // Line 4
        }
    

  • 0
    B

    I find the second iterative solution quite hard to understand for someone (like me) who doesn't know much Python, so I'm really not sure if my solution below has a lot of (or any?) similarities to the second solution. It basically is a standard iterative post-order traversal, with the path checking for nodes p and q. But I hope it's easy to understand!

    C++ Post-Order Iterative Traversal Solution:

    TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
            TreeNode *cur=root,*last=nullptr;
            vector<TreeNode*> pathp,pathq,temp;
            
            while (pathp.empty() || pathq.empty()) {
                
                // standard post-order iterative tree traversal
                if (cur) {
                    temp.push_back(cur);
                    if (temp.back()==p) pathp=temp; // check and set path for p
                    if (temp.back()==q) pathq=temp; // check and set path for q
                    cur=cur->left;
                } else {
                    if (temp.back()->right && temp.back()->right != last) {
                        cur=temp.back()->right;
                    } else {
                        last=temp.back();
                        temp.pop_back();
                    }
                }
            }
            
            // compare paths and get lowest common ancestor
            int n=min(pathp.size(),pathq.size());
            for (int i=1; i<n; i++) {
                if (pathp[i]!=pathq[i]) return pathp[i-1];
            }
            return pathp[n-1];
        }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.