Data-Structure/模板/树/两序重建树.cpp

#include <iostream>
#include <vector>
#include <unordered_map>

using namespace std;

/*====================================================
=                   二叉树节点定义                    =
====================================================*/
struct TreeNode {
    int val;
    TreeNode* left;
    TreeNode* right;

    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};

/*====================================================
=          前序 + 中序 重建二叉树（模块一）            =
====================================================*/
class BuildFromPreIn {
public:
    unordered_map<int, int> pos;   // 中序遍历中：值 -> 下标

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

private:
    TreeNode* dfs(vector<int>& pre, int preL, int preR,
                  vector<int>& in,  int inL,  int inR) {
        if (preL > preR) return nullptr;

        // 前序第一个元素是当前子树的根
        int rootVal = pre[preL];
        TreeNode* root = new TreeNode(rootVal);

        // 根在中序中的位置
        int k = pos[rootVal];
        int leftSize = k - inL;

        // 构建左子树
        root->left = dfs(
            pre,
            preL + 1,
            preL + leftSize,
            in,
            inL,
            k - 1
        );

        // 构建右子树
        root->right = dfs(
            pre,
            preL + leftSize + 1,
            preR,
            in,
            k + 1,
            inR
        );

        return root;
    }
};

/*====================================================
=          中序 + 后序 重建二叉树（模块二）            =
====================================================*/
class BuildFromInPost {
public:
    unordered_map<int, int> pos;   // 中序遍历中：值 -> 下标

    TreeNode* buildTree(vector<int>& inorder, vector<int>& postorder) {
        int n = inorder.size();
        pos.clear();
        for (int i = 0; i < n; i++) {
            pos[inorder[i]] = i;
        }
        return dfs(inorder, 0, n - 1, postorder, 0, n - 1);
    }

private:
    TreeNode* dfs(vector<int>& in,  int inL,   int inR,
                  vector<int>& post,int postL, int postR) {
        if (inL > inR) return nullptr;

        // 后序最后一个元素是当前子树的根
        int rootVal = post[postR];
        TreeNode* root = new TreeNode(rootVal);

        // 根在中序中的位置
        int k = pos[rootVal];
        int leftSize = k - inL;

        // 构建左子树
        root->left = dfs(
            in,
            inL,
            k - 1,
            post,
            postL,
            postL + leftSize - 1
        );

        // 构建右子树
        root->right = dfs(
            in,
            k + 1,
            inR,
            post,
            postL + leftSize,
            postR - 1
        );

        return root;
    }
};

/*====================================================
=              辅助函数：遍历与验证                   =
====================================================*/

// 中序遍历（用于验证结构是否正确）
void inorderPrint(TreeNode* root) {
    if (!root) return;
    inorderPrint(root->left);
    cout << root->val << ' ';
    inorderPrint(root->right);
}

// 后序遍历
void postorderPrint(TreeNode* root) {
    if (!root) return;
    postorderPrint(root->left);
    postorderPrint(root->right);
    cout << root->val << ' ';
}

// 前序遍历
void preorderPrint(TreeNode* root) {
    if (!root) return;
    cout << root->val << ' ';
    preorderPrint(root->left);
    preorderPrint(root->right);
}

/*====================================================
=                        main                         =
====================================================*/
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    /*
        示例测试数据
        前序:  3 9 20 15 7
        中序:  9 3 15 20 7
        后序:  9 15 7 20 3
    */

    vector<int> preorder  = {3, 9, 20, 15, 7};
    vector<int> inorder   = {9, 3, 15, 20, 7};
    vector<int> postorder = {9, 15, 7, 20, 3};

    // 前序 + 中序
    BuildFromPreIn solver1;
    TreeNode* root1 = solver1.buildTree(preorder, inorder);

    // 中序 + 后序
    BuildFromInPost solver2;
    TreeNode* root2 = solver2.buildTree(inorder, postorder);

    // 验证输出
    inorderPrint(root1);
    cout << '\n';

    inorderPrint(root2);
    cout << '\n';

    return 0;
}