剑指7：重建二叉树

传送门

nowcoder
leetcode

题目描述

输入某二叉树的前序遍历和中序遍历的结果，请重建出该二叉树。
假设输入的前序遍历和中序遍历的结果中都不含重复的数字。
例如：输入前序遍历序列 {1,2,4,7,3,5,6,8} 和中序遍历序列 {4,7,2,1,5,3,8,6}，则重建二叉树并返回。

思路

根据二叉树先序遍历与中序遍历的规则。
前序遍历的起始元素为树的根节点 node 的值。
在中序遍历中搜索根节点 node 的索引，将中序遍历划分为 [左子树 | 根节点 | 右子树]
根据中序遍历中的左/右子树的节点数量，将前序遍历划分为 [根节点 | 左子树 | 右子树]

pre 的起始元素作为根节点，在 inorder 找到根节点的索引 mid
pre[1, mid] 为左子树，pre[mid+1, length] 为右子树
inorder[0, mid-1] 为左子树，inorder[mid+1, length] 为右子树

C++ 代码 - nowcoder

class Solution {
public:
    TreeNode* reConstructBinaryTree(vector<int>& pre, vector<int>& in) {
        unordered_map<int, int> num2index;
        for (int i = 0; i < in.size(); i ++) {
            num2index.insert({in[i], i});
        }
        return reConstructBinaryTreeCore(num2index, 
            pre, 0, pre.size() - 1,
            in, 0, in.size() - 1);
    }

    TreeNode* reConstructBinaryTreeCore(unordered_map<int, int> num2index,
        vector<int>& pre, int pl, int pr,
        vector<int>& in, int il, int ir) {
        if (pl > pr || il > ir) return nullptr;

        TreeNode* root = new TreeNode(pre[pl]);
        int in_root_index = num2index[root->val];
        root->left = reConstructBinaryTreeCore(num2index,
            pre, pl + 1, pl + in_root_index - il,
            in, il, in_root_index - 1);
        root->right = reConstructBinaryTreeCore(num2index,
            pre, pl + in_root_index - il + 1, pr,
            in, in_root_index + 1, ir);
        return root;
    }
};

/*
    C++ STL distance()
*/
class Solution {
public:
    TreeNode* reConstructBinaryTree(vector<int>& pre, vector<int>& in) {
        if (pre.size() == 0 || in.size() == 0) return nullptr;

        TreeNode* rootNode = new TreeNode(pre[0]);
        int mid = distance(begin(in), find(in.begin(), in.end(), pre[0]));
        vector<int> pre_left(pre.begin() + 1, pre.begin() + mid + 1);
        vector<int> pre_right(pre.begin() + mid + 1, pre.end());
        vector<int> in_left(in.begin(), in.begin() + mid);
        vector<int> in_right(in.begin() + mid + 1, in.end());

        rootNode->left = reConstructBinaryTree(pre_left, in_left);
        rootNode->right = reConstructBinaryTree(pre_right, in_right);
        return rootNode;
    }
};

C++ 代码 - leetcode

class Solution {
public:
    unordered_map<int, int> num2idx;

    TreeNode* buildTree(vector<int>& preorder, vector<int>& inorder) {
        for (int i = 0; i < inorder.size(); i ++) {
            num2idx.insert({inorder[i], i});
        }
        return build(preorder, 0, preorder.size() - 1, inorder, 0, inorder.size() - 1);
    }

    /* 
        pl, pr 表示子树的前序遍历在数组中的位置 [left, right]
        il, ir 表示子树的中序遍历在数组中的位置 [left, right]
    */
    TreeNode* build(vector<int>& preorder, int pl, int pr,
                    vector<int>& inorder, int il, int ir) {
        if (pl > pr || il > ir) return nullptr;

        TreeNode* rootNode = new TreeNode(preorder[pl]);
        int inRoot = num2idx[rootNode->val];
        int leftSubTreeNodes = inRoot - il;
        rootNode->left = build(preorder, pl + 1, pl + leftSubTreeNodes, inorder, il, inRoot - 1);
        rootNode->right = build(preorder, pl + leftSubTreeNodes + 1, pr, inorder, inRoot + 1, pr);
        return rootNode;
    }
};