An extremly straightforward way to write iterative code from recursions (for dummies)


  • 11
    O

    Here I provide a very straightforward and general way to transform recursion into iterative code. This works especially well for dummies like me.

    As we know, the recursion code is very easy to write

        void recursion(TreeNode * cur, vector<int> & result) {
            if (cur->left) recursion(cur->left);        // line 0
            if (cur->right) recursion(cur->right);      // line 1
            result.push_back(cur->val);                 // line 2
            return;                                     // line 3
        }
    

    The difficulty is to write the iterative version of such code. I fully mimic the function call stacks. I maintain two stacks: variable stack and pc stack (pc stands for programming counter). Here the variable stack maintains the value of each variable of each function call in the recursion. The pc stack stores the line of code that is to be executed in each corresponding function

    The entry point in the recursion code is to call recursion(root) function. Hence, in the iterative version of the code, we push root to the variable stack, and 0 to the pc stack. This means that the first line (line 0) of recursion(root) function is to be executed.

    Then I execute the following loop until the stack is empty.

    1. Inspect the top of the pc stack. Then increment the top of pc stack, which means that the current line of code has been executed.
    2. If pc == 0, push the left child of the top of the variable stack to the stack if such child exists. Also push 0 to the pc stack.
    3. If pc == 1, push the right child of the top of the variable stack to the stack if such child exists. Also push 0 to the pc stack.
    4. if pc == 2, append the value of the top of the variable stack to the result.
    5. If pc == 3, pop both variable and pc stack.

    You can find pc here fully corresponds to the line number in the previous recursion code!

    The complete code is following.

    class Solution {
    public:
        vector<int> postorderTraversal(TreeNode* root) {
            stack<TreeNode *> varStack;
            stack<int> PCStack;
            vector<int> result;
            if (!root) return result;
            varStack.push(root);
            PCStack.push(0);
            while (!varStack.empty()) {
                TreeNode * cur = varStack.top();
                int pc = PCStack.top();
                PCStack.top()++;
                if (pc == 0) {
                    // line 0: if (cur->left) recursion(cur->left); 
                    if (cur->left) {
                        varStack.push(cur->left);
                        PCStack.push(0);
                    }
                }
                else if (pc == 1) {
                    // line 1: if (cur->right) recursion(cur->right); 
                    if (cur->right) {
                        varStack.push(cur->right);
                        PCStack.push(0);
                    }
                }
                else if (pc == 2) {
                    // line 2: result.push_back(cur->val);
                    result.push_back(cur->val);
                }
                else if (pc == 3) {
                    // line 3: return
                    varStack.pop();
                    PCStack.pop();
                }
            }
            return result;
        }
    };
    

    Yes it is LONG. But it is extremely easy to write, since you are simply simulating how the computer does for recursion.

    The most amazing part of such method is that, it works for all recursions!!! I think that it also shows your understanding of call stacks if you analyze recursion in such a straightforward way.


  • 0
    P

    Thank you so much!! Wonderful insight into converting recursive solution to iterative. :)


  • 0
    C

    @oldfish

    I has the same solution with you

    /**
     * Definition for a binary tree node.
     * struct TreeNode {
     *     int val;
     *     struct TreeNode *left;
     *     struct TreeNode *right;
     * };
     */
    typedef struct stack {
        void           *c;
        struct stack   *next;
    } Stack;
    
    void create(Stack **ptop)
    {
        *ptop = NULL;
    }
    
    void push(Stack **ptop, void *c)
    {
        Stack   *tmp;
    
        tmp = malloc(sizeof *tmp);
        tmp->c = c;
        tmp->next = *ptop;
    
        *ptop = tmp;
    }
    
    void pop(Stack **ptop, void **c)
    {
        Stack   *tmp;
    
        if (! *ptop) return;
    
        tmp = *ptop;
        if (c) *c = tmp->c;
        *ptop = (*ptop)->next;
        free(tmp);
    }
    
    void top(Stack **ptop, void **c)
    {
        *c = (*ptop)->c;
    }
    
    int empty(Stack **ptop)
    {
       return *ptop == NULL; 
    }
    
    /**
     * Return an array of size *returnSize.
     * Note: The returned array must be malloced, assume caller calls free().
     */
    int* postorderTraversal(struct TreeNode* root, int* returnSize)
    {
        int             *arr;
        long            mark;
        Stack           *s;
        Stack           *pc;
        
        create(&s);
        create(&pc);
        arr = malloc(10000 * sizeof *arr);
        *returnSize = 0;
        
        push(&s, root);
        push(&pc, 0);
        while (!empty(&s)) {
            top(&s, &root);
    
            if (!root) {
                pop(&s, &root);
                pop(&pc, &mark);
            } else {
                pop(&pc, &mark);
                
                if (mark == 0) {
                    push(&s, root->left);
                    push(&pc, 1);
                    push(&pc, 0);
                } else if (mark == 1) {
                    push(&s, root->right);
                    push(&pc, 2);
                    push(&pc, 0);
                } else {
                    pop(&s, &root);
                    arr[(*returnSize)++] = root->val;
                }
            }
        }
        
        return arr;
    }
    

Log in to reply
 

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