二叉树计算 - 华为OD统一考试（C卷）

OD统一考试（C卷）

分值： 200分

题解： Java / Python / C++

alt

题目描述

给出一个二叉树如下图所示：

请由该二叉树生成一个新的二叉树，它满足其树中的每个节点将包含原始树中的左子树和右子树的和。

      20 (7-2+9+6)
     /   \
    -2    6
     \   /  
     0  0

左子树表示该节点左侧叶子节点为根节点的一颗新树；右子树表示该节点右侧叶子节点为根节点的一颗新树

输入描述

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

例如：

7 -2 6 6 9 6 7 -2 9 6

输出描述

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

例如：

输出1 -2 0 20 0 6

示例1

输入：
-3 12 6 8 9 -10 -7
8 12 -3 6 -10 9 -7

输出：
0 3 0 7 0 2 0

题解

递归的题目

根据树的中序和先序遍历可以唯一确定一颗树，题解中将中序和先序字符串看成一个整体（树）。

题目求的是求和树的中序遍历，因此还是使用递归中序遍历的方式去遍历树，遍历的同时去求左右子树和（即为答案）。

Java

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;
import java.util.stream.Collectors;

/**
 * @author code5bug
 */
public class Main {
    static List<Integer> rs = new ArrayList<>();

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);

        // 中序遍历
        List<Integer> inorder = Arrays.stream(scanner.nextLine().split(" "))
                .map(p -> Integer.parseInt(p)).collect(Collectors.toList());
        // 先序遍历
        List<Integer> preorder = Arrays.stream(scanner.nextLine().split(" "))
                .map(p -> Integer.parseInt(p)).collect(Collectors.toList());

        solve(inorder, preorder);

        for (int num : rs) {
            System.out.print(num + " ");
        }
    }

    static void solve(List<Integer> inorder, List<Integer> preorder) {
        if (inorder.isEmpty() || preorder.isEmpty()) {
            return;
        }

        // 数根节点的值
        int rootVal = preorder.get(0);
        int idx = inorder.indexOf(rootVal);

        // 中序，左右子树
        List<Integer> leftInorder = inorder.subList(0, idx);
        List<Integer> rightInorder = inorder.subList(idx + 1, inorder.size());

        // 先序，左右子树
        List<Integer> leftPreorder = preorder.subList(1, idx + 1);
        List<Integer> rightPreorder = preorder.subList(idx + 1, preorder.size());

        solve(leftInorder, leftPreorder);

        // 左右子树求值求和
        int leftSum = leftInorder.stream().reduce(0, Integer::sum);
        int rightSum = rightInorder.stream().reduce(0, Integer::sum);
        rs.add(leftSum + rightSum);

        solve(rightInorder, rightPreorder);
    }

}

Python

# 以下 2 行代码解决递归深度报错只能 AC 95% 的问题
import sys
sys.setrecursionlimit(10000)

# 中序遍历
inorder = list(map(int, input().split()))
# 先序遍历
preorder = list(map(int, input().split()))

# 求和树的中序遍历
rs = []

def solve(inorder, preorder):
    if not inorder or not preorder:
        return
    
    root_val = preorder[0]
    idx = inorder.index(root_val)

    left_inorder = inorder[: idx]
    right_inorder = inorder[idx + 1:]

    left_preorder = preorder[1: idx + 1]
    right_preorder = preorder[idx + 1:]

    solve(left_inorder, left_preorder)
    rs.append(sum(left_inorder) + sum(right_inorder))
    solve(right_inorder, right_preorder)

solve(inorder, preorder)
print(*rs)

C++

#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>

using namespace std;

vector<int> rs;

void solve(vector<int>& inorder, vector<int>& preorder) {
    if (inorder.empty() || preorder.empty()) {
        return;
    }

    // 数根节点的值
    int rootVal = preorder[0];
    auto it = find(inorder.begin(), inorder.end(), rootVal);
    int idx = distance(inorder.begin(), it);

    // 中序，左右子树
    vector<int> leftInorder(inorder.begin(), inorder.begin() + idx);
    vector<int> rightInorder(inorder.begin() + idx + 1, inorder.end());

    // 先序，左右子树
    vector<int> leftPreorder(preorder.begin() + 1, preorder.begin() + idx + 1);
    vector<int> rightPreorder(preorder.begin() + idx + 1, preorder.end());

    solve(leftInorder, leftPreorder);

    // 左右子树求值求和
    int leftSum = accumulate(leftInorder.begin(), leftInorder.end(), 0);
    int rightSum = accumulate(rightInorder.begin(), rightInorder.end(), 0);
    rs.push_back(leftSum + rightSum);

    solve(rightInorder, rightPreorder);
}

vector<int> read_line(vector<int>& collect) {
    int v;
    while (cin >> v) {
        collect.push_back(v);
        if (cin.peek() == '\n') break;
    }
    return collect;
}

int main() {
    // 中序遍历
    vector<int> inorder;
    read_line(inorder);

    // 先序遍历
    vector<int> preorder;
    read_line(preorder);

    solve(inorder, preorder);

    for (int num : rs) {
        cout << num << " ";
    }

    return 0;
}