[leetcode] 5031. 从先序遍历还原二叉树

题目

我们从二叉树的根节点 root 开始进行深度优先搜索。

在遍历中的每个节点处，我们输出 D 条短划线（其中 D 是该节点的深度），然后输出该节点的值。（如果节点的深度为 D，则其直接子节点的深度为 D + 1。根节点的深度为 0）。

如果节点只有一个子节点，那么保证该子节点为左子节点。

给出遍历输出 S，还原树并返回其根节点 root。

p.s. 示例有图片有点麻烦所以给原链接 https://leetcode-cn.com/contest/weekly-contest-132/problems/recover-a-tree-from-preorder-traversal/

解题思路

lc132场周赛的最后一题。这题的数据结构是树，是把dfs遍历的字符串还原成原始结构的过程。

考虑到dfs的特性应该使用stack结构。使用一个stack<pair<TreeNode*, int>>记录树节点和对应的层数。

核心思路：取得新的pair，观察stack状态

（1）如果pair的层数和stack顶层数相同，说明互为兄弟节点，入栈

（2）如果pair的层数比stack顶层数大，说明为后代节点，入栈

（3）如果pair的层数比stack顶层数小，说明在该层之前的树已经遍历完毕，可以合并为一棵树，出栈（stackCompaction）

（4）如果栈为空，说明为root节点，入栈

（5）重复

最后栈里只会剩一个节点就是建好的root节点

 

代码

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    int getVal(int& p, string S){
        int val = 0;
        while(p < S.size() && S[p] >= '0' && S[p] <= '9'){
            val = val * 10 + S[p] - '0';
            p += 1;
        }
        return val;
    }
    
    int getDep(int& p, string S){
        int d = 0;
        while(p < S.size() && S[p] == '-'){
            d += 1;
            p += 1;
        }
        return d;
    }
    
    void stackCompaction(stack<pair<TreeNode*, int>>& sta){
        int rDep = sta.top().second;
        TreeNode* r = sta.top().first;
        sta.pop();
        //empty impossible unless rt
        int lDep = sta.top().second;
        TreeNode* l = sta.top().first;
        sta.pop();
        if(rDep == lDep){
            int pDep = sta.top().second;
            TreeNode* pn = sta.top().first;
            sta.pop();
            pn->left = l;
            pn->right = r;
            pair<TreeNode*, int> newParent = make_pair(pn, pDep);
            sta.push(newParent);
        }
        else if(rDep > lDep){ //only one child
            TreeNode* pn = l;
            int pDep = lDep;
            pn -> left = r;
            pair<TreeNode*, int> newParent = make_pair(pn, pDep);
            sta.push(newParent);
        }
        else{//impossible
        }
    }
    
    TreeNode* recoverFromPreorder(string S) {
        int p = 0;
        stack<pair<TreeNode*, int>> sta; //node, dep
        while(p < S.size()){
            int dep = getDep(p, S);
            int val = getVal(p, S);

            TreeNode* n = new TreeNode(val);
            pair<TreeNode*, int> i = make_pair(n, dep);
            
            if(sta.empty() || sta.top().second <= dep){
                sta.push(i);
            }
            else if(sta.top().second > dep){
                while(sta.size() >= 2 && sta.top().second > dep){
                    stackCompaction(sta);
                }
                sta.push(i);
            }
        }
        
        while(sta.size() >= 2){
            stackCompaction(sta);
        }
        return sta.top().first;
    }
};