[LeetCode] 1028. Recover a Tree From Preorder Traversal

https://leetcode.com/problems/recover-a-tree-from-preorder-traversal/description/

leetcode - Recover a Tree From Preorder Traversal

문제 유형: 스택, 트리, dfs, 문자열 처리

문제 난이도: Hard

 

문제

We run a preorder depth-first search (DFS) on the root of a binary tree.

At each node in this traversal, we output D dashes (where D is the depth of this node), then we output the value of this node.  If the depth of a node is D, the depth of its immediate child is D + 1.  The depth of the root node is 0.

If a node has only one child, that child is guaranteed to be the left child.

Given the output traversal of this traversal, recover the tree and return its root.

 

이진 트리의 프리오더 순회가 주어집니다.

각 노드의 순회에서, 깊이D만큼의 대시가 붙습니다.

루트의 깊이는 0입니다.

만약 노드가 하나의 자식만 있다면, 그 자식은 반드시 왼쪽 자식입니다.

순회결과가 문자열로 주어지면, 이 정보를 이용해서 복구해서 루트를 반환하시오.

 

풀이

우선, {노드 값, 레벨}의 꼴로 문자열을 파싱해서 노드정보 배열을 만든다.

투 포인터를 이용해서 쉽게 할 수 있다.

 

스택을 만든다. <노드 주소, 레벨>을 저장할 것이다.

그 뒤, 스택에 루트를 삽입하고, 다음을 노드정보 배열을 한 번 순회하는 동안 진행한다:

  왼쪽우선으로 계속 더 높은 레벨을 만날때마다 왼쪽자식으로 추가하고, 스택에 넣는다.

  그 뒤, 이제 오른쪽을 추가해야 한다.

  부모 노드(이번에 추가할 노드보다 작은 레벨의 최초 노드)로 스택에서 백트래킹한다.

  그 노드에서 오른쪽자식을 추가하고, 스택에 그 자식을 넣는다.

이를 계속 반복하면, 완성된다.

스택을 이용해서 preorder traversal한 것과 비슷하다고 보면 된다.

 

코드

C++

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    TreeNode* recoverFromPreorder(string traversal) {
        int n = traversal.size();
        vector<pair<int, int>> nodes; // val, level
        int l = 0;
        int r = 0;
        int level = 0;
        while(l < n) {
            // node value
            while(r < n && traversal[r] >= '0' && traversal[r] <= '9') {
                r++;
            }

            int v = stoi(traversal.substr(l, r-l));
            nodes.push_back({v, level});
            
            // level 
            l = r;
            while(r < n && traversal[r] == '-') {
                r++;
            }
            level = r-l;

            l = r;
        }

        TreeNode* root = new TreeNode(nodes[0].first);
        stack<pair<TreeNode*, int>> stk;
        stk.push({root, 0});
        int i = 1;
        int m = nodes.size();
        while(i < m) {
            while(i < m && stk.top().second < nodes[i].second) {
                auto curr = stk.top();
                auto currNode = curr.first;
                auto currLevel = curr.second;

                auto newNode = new TreeNode(nodes[i].first);
                currNode -> left = newNode;
                stk.push({newNode, nodes[i].second});
                i++;
            }
            
            if(i < m) {
                while(!stk.empty() && stk.top().second >= nodes[i].second) {
                    stk.pop();
                }

                auto currNode = stk.top().first;
                auto newNode = new TreeNode(nodes[i].first);
                currNode -> right = newNode;
                stk.push({newNode, nodes[i].second});
                i++;
            }
        }

        return root;
    }
};

 

Go

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func recoverFromPreorder(traversal string) *TreeNode {
    n := len(traversal)
    l := 0
    r := 0
    level := 0
    nodes := []struct{
        val int
        level int
    }{}

    for l < n {
        for r < n && traversal[r] >= '0' && traversal[r] <= '9' {
            r++
        }
        v, _ := strconv.Atoi(traversal[l:r])
        nodes = append(nodes, struct{
            val int
            level int
        }{v, level})
        l = r

        for r < n && traversal[r] == '-' {
            r++
        }
        level = r-l
        l = r 
    }

    root := &TreeNode{Val: nodes[0].val}
    stack := []struct{
        node *TreeNode
        level int
    }{{root, 0}}
    i := 1
    m := len(nodes)

    for i < m {
        for i < m && stack[len(stack)-1].level < nodes[i].level {
            top := stack[len(stack)-1]
            newNode := &TreeNode{Val: nodes[i].val}
            top.node.Left = newNode
            stack = append(stack, struct {
                node *TreeNode
                level int
            }{newNode, nodes[i].level})
            i++
        }

        if i < m {
            for len(stack) > 0 && stack[len(stack)-1].level >= nodes[i].level {
                stack = stack[:len(stack)-1]
            }

            parent := stack[len(stack)-1].node
            newNode := &TreeNode{Val: nodes[i].val}
            parent.Right = newNode
            stack = append(stack, struct{
                node *TreeNode
                level int
            }{newNode, nodes[i].level})
            i++
        }
    }

    return root
}