重建树

import bean.TreeNode;

import java.util.*;

public class TreeTest {

    public static final int PRE_ORDER = 1;
    public static final int MID_ORDER = 2;
    public static final int AFTER_ORDER = 3;

    static List preOrder = new ArrayList<>();
    static List midOrder = new ArrayList<>();
    static List afterOrder = new ArrayList<>();

    public static void main(String[] args) {

        TreeNode root = new TreeNode(3);
        root.left = new TreeNode(6);
        root.right = new TreeNode(5);
        root.left.right = new TreeNode(9);
        root.right.left = new TreeNode(10);
        root.right.right = new TreeNode(20);
        new TreeTest().traversal(root);
        int[] preArray = new int[preOrder.size()];
        for (int i = 0; i < preOrder.size(); i++) {
            preArray[i] = preOrder.get(i);
            System.out.print(" " + preOrder.get(i));
        }
        System.out.println();
        int[] midArray = new int[midOrder.size()];
        for (int i = 0; i < midOrder.size(); i++) {
            midArray[i] = midOrder.get(i);
            System.out.print(" " + midOrder.get(i));
        }
        System.out.println();
        for (int i = 0; i < afterOrder.size(); i++) {
            System.out.print(" " + afterOrder.get(i));
        }

        
        // 根据先序和中序重建树
        TreeNode rebuild = new TreeNode(preArray[0]);
        Map map = new HashMap<>();
        for (int i = 0; i < midArray.length; i++) {
            map.put(midArray[i], i);
        }
        Stack stack = new Stack<>();
        stack.push(rebuild);
        for (int i = 1; i < preArray.length; i++) {
            TreeNode top = stack.peek();
            // 栈顶值在mid中的位置
            int p = map.get(top.val);
            //插入值在mid中的位置
            int j = map.get(preArray[i]);

            if (j < p) {
                top.left = new TreeNode(preArray[i]);
                stack.push(top.left);
            } else {
                TreeNode q = null;
                while (!stack.isEmpty() && map.get(stack.peek().val) < j) {
                    q = stack.pop();
                }

                q.right = new TreeNode(preArray[i]);
                stack.push(q.right);
            }
        }

        System.out.println(isSameTree(root, rebuild));

    }

    // 判断两棵树是否相同
    public static boolean isSameTree(TreeNode a, TreeNode b) {
        if (a == null && b == null) {
            return true;
        }

        if (a == null || b == null) {
            return false;
        }

        return a.val == b.val && isSameTree(a.left, b.left) && isSameTree(a.right, b.right);
    }

    // 树的遍历
    public void traversal(TreeNode root) {
        if (root == null) {
            return;
        }
        preOrder.add(root.val);
        traversal(root.left);
        midOrder.add(root.val);
        traversal(root.right);
        afterOrder.add(root.val);

    }

}