【Tree】已知前序和中序遍历还原二叉树

/**************************
https://leetcode.com/problems/construct-binary-tree-from-preorder-and-inorder-traversal/
@date 2015.5.16
@description
根据前序和中序遍历，重构二叉树
可用迭代，也可以使用递归。
递归前面已经写过，耗时，这里再次做使用迭代。

@tags tree, array, depth-first search
@solution
这里tag里有DFS，迭代和递归都用到。
参照先序遍历的迭代的思想，总是把左子树的左节点优先push入栈。
当栈顶元素s.top() == inorder[0]时，则说明已到达树的最左的叶子节点，随后按照先序遍历迭代思想，转向右子树。
这里设置了一个flag标志位，flag = 1转向右子树，默认flag=0持续将左子树左节点入栈。
***************************/
#include <iostream>
#include <vector>
#include <stack>

using namespace std;

struct TreeNode{
    int val;
    TreeNode *left, *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution {
public:
    TreeNode* buildTree(vector<int>& preorder, vector<int>& inorder){
        if (preorder.size() == 0)
            return NULL;

        stack<TreeNode *> s;

        int i = 0; // 指向先序遍历向量
        int j = 0; // 指向中序遍历向量
        int flag = 0; // flag = 0指向左子树，当flag = 1指向右子树

        TreeNode *root = new TreeNode(preorder[i]);
        s.push(root);
        TreeNode *temp = root;

        i++;

        while (i < inorder.size()){
            if (!s.empty() && s.top()->val == inorder[j]){
                temp = s.top(); // 重要，更换temp
                s.pop();
                flag = 1;
                j++;
            } else{
                if (flag == 0){
                    temp->left = new TreeNode(preorder[i]);
                    temp = temp->left;
                    s.push(temp);
                    i++;
                } else{
                    flag = 0; // 右节点不连续添加
                    temp->right = new TreeNode(preorder[i]);
                    temp = temp->right;
                    s.push(temp);
                    i++;
                }
            }
        }
        return root;

    }


};

int main(){
    vector<int> preorder = {1, 2, 4, 5, 3, 6, 7};
    vector<int> inorder = {4, 2, 5, 1, 6, 3, 7};
    TreeNode * root;
    Solution a;
    root = a.buildTree(preorder, inorder);
    cout << root->val << endl;
    cout << root->left->val << endl;
    cout << root->right->val << endl;
}