LeetCode 每日一题 11028. 从先序遍历还原二叉树

题目描述

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

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

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

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

输入："1-2--3--4-5--6--7"
输出：[1,2,5,3,4,6,7]

输入："1-2--3---4-5--6---7"
输出：[1,2,5,3,null,6,null,4,null,7]

输入："1-401--349---90--88"
输出：[1,401,null,349,88,90]

提示：

• 原始树中的节点数介于 1 和 1000 之间。
• 每个节点的值介于 1 和 10 ^ 9 之间。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/recover-a-tree-from-preorder-traversal
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

My Solution

DFS 用栈来保存路径，这里用的 stack<pair<TreeNode*, int>>stk; 同时记录节点的深度。
DFS 时显然有深度是递增的。利用这一点找到当前节点的父节点即可。

#pragma G++ optimize("O2")

/**
 * 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:
  pair<int, int> parseString(string &s, int& idx) {
    int depth = 0;
    while(s[idx] == '-')
      depth++, idx++;
    int val = atoi(s.c_str() + idx);
    while(isdigit(s[idx]))
      idx++;
    //cout << ": " << val << " " << depth << endl;
    return make_pair(val, depth);
  }

  TreeNode* recoverFromPreorder(string S) {
    stack<pair<TreeNode*, int>>stk;
    pair<TreeNode*, int>u;
    TreeNode*root = nullptr;
    int idx = 0;
    while(idx < S.length()) {
      const pair<int, int> u = parseString(S, idx);
      if(stk.size() == 0) {
        root = new TreeNode(u.first);
        stk.push(make_pair(root, u.second));
      } else {
        while(stk.top().second >= u.second)
          stk.pop();

        TreeNode *node = new TreeNode(u.first);
        if(stk.top().first->left == nullptr) {
          stk.top().first->left = node;
          stk.push(make_pair(node, u.second));
        } else if(stk.top().first->right == nullptr) {
          stk.top().first->right = node;
          stk.push(make_pair(node, u.second));
        } else {
          cout << " ?? " << endl;
        }
      }
    }
    return root;
  }
};