二叉树计算 - 技术栈

题目描述

给出一个二叉树，请由该二叉树生成一个新的二叉树，它满足其树中的每个节点将包含原始树中的左子树和右子树的和。左子树表示该节点左侧叶子节点为根节点的一颗新树;右子树表示该节点右侧叶子节点为根节点的一颗新树。

输入描述

2行整数，第1行表示二叉树的中序遍历，第2行表示二叉树的前序遍历，以空格分割。

输出描述

1行整数，表示求和树的中序遍历，以空格分割。

例1：

复制代码

输入：
-3 12 6 8 9 -10 -7
8 12 -3 6 -10 9 -7
输出：
0 3 0 7 0 2 0

java 复制代码

/*
-3 12 6 8 9 -10 -7
8 12 -3 6 -10 9 -7
0 3 0 7 0 2 0
 */
public class 二叉树计算 {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int[] mid = Arrays.stream(sc.nextLine().trim().split(" ")).mapToInt(Integer::parseInt).toArray();
        int[] pre = Arrays.stream(sc.nextLine().trim().split(" ")).mapToInt(Integer::parseInt).toArray();
        // 构建树
        Node root = buildTree(mid, pre);
        // 计算每个节点的值
        sumTree(root);
        // 中序遍历输出结果
        printRes(root);
    }

    private static void printRes(Node root) {
        if (root == null){
            return;
        }
        printRes(root.left);
        System.out.print(root.val + " ");
        printRes(root.right);
    }

    private static Integer sumTree(Node node) {
        if (node == null){
            return 0;
        }
        int nodeLeftSum = sumTree(node.left);
        int nodeRightSum = sumTree(node.right);
        int valOld = node.val;
        node.val = nodeLeftSum + nodeRightSum;
        return node.val + valOld;
    }

    private static Node buildTree(int[] mid, int[] pre) {
        HashMap<Integer, Integer> midMap = new HashMap<>();
        for (int i = 0; i < mid.length; i++) {
            midMap.put(mid[i], i);
        }
        return getTree(pre, 0, pre.length-1, mid, 0, mid.length-1, midMap);
    }

    private static Node getTree(int[] pre, int preIndexStart, int preIndexEnd, int[] mid,
                                int midIndexStart, int midIndexend, HashMap<Integer, Integer> midMap) {
        if (preIndexStart > preIndexEnd || midIndexStart > midIndexend){
            return null;
        }

        int rootVal = pre[preIndexStart];
        Node root = new Node(rootVal);
        // 根据root节点在中序遍历中的下标，可以获取root节点的左右节点的长度
        Integer midRootIndex = midMap.get(rootVal);
        int leftSize = midRootIndex - midIndexStart;

        root.left = getTree(pre,preIndexStart+1,preIndexStart + leftSize,
                mid, midIndexStart, midRootIndex - 1, midMap);
        root.right = getTree(pre,preIndexStart + leftSize + 1,preIndexEnd,
                mid, midRootIndex + 1, midIndexend, midMap);

    return root;
    }

    static class Node{
        int val;
        Node left;
        Node right;

        public Node(int val) {
            this.val = val;
        }
    }
}