[C++知识库]PTA甲级 1102 Invert a Binary Tree (C++)

The following is from Max Howell @twitter:

Google: 90% of our engineers use the software you wrote (Homebrew), but you can’t invert a binary tree on a whiteboard so fuck off.

Now it’s your turn to prove that YOU CAN invert a binary tree!

Input Specification:

Each input file contains one test case. For each case, the first line gives a positive integer $N(\le 10)$ which is the total number of nodes in the tree – and hence the nodes are numbered from 0 to $N?1$. Then $N$ lines follow, each corresponds to a node from 0 to $N?1$, and gives the indices of the left and right children of the node. If the child does not exist, a - will be put at the position. Any pair of children are separated by a space.

Output Specification:

For each test case, print in the first line the level-order, and then in the second line the in-order traversal sequences of the inverted tree. There must be exactly one space between any adjacent numbers, and no extra space at the end of the line.

Sample Input:

8
1 -
- -
0 -
2 7
- -
- -
5 -
4 6

Sample Output:

3 7 2 6 4 0 5 1
6 5 7 4 3 2 0 1

Solution:

// Talk is cheap, show me the code
// Created by Misdirection 2021-08-24 21:30:09
// All rights reserved.

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

using namespace std;

vector<pair<int, int>> input;
vector<int> levelOrder, midOrder;

struct Node{
    int key;
    Node *left;
    Node *right;

    Node(int k = -1, Node *l = NULL, Node *r = NULL){
        key = k;
        left = l;
        right = r;
    }
    ~Node(){}
};

void findChildren(Node *root){
    if(root == NULL) return;

    if(input[root -> key].first != -1) root -> left = new Node(input[root -> key].first, NULL, NULL);
    if(input[root -> key].second != -1) root -> right = new Node(input[root -> key].second, NULL, NULL);

    findChildren(root -> left);
    findChildren(root -> right);
}

void midTraverse(Node *root){
    if(root -> left != NULL) midTraverse(root -> left);
    midOrder.push_back(root -> key);
    if(root -> right != NULL) midTraverse(root -> right);
}

int main(){
    int n;
    scanf("%d", &n);

    int r;
    unordered_map<int, int> flag;
    input.clear();
    levelOrder.clear();
    midOrder.clear();
    
    for(int i = 0; i < n; ++i){
        string a, b;
        cin >> a >> b;

        int left, right;

        if(a == "-") left = -1;
        else left = stoi(a);

        if(b == "-") right = -1;
        else right = stoi(b);

        flag[left] = flag[right] = 1;
        input.emplace_back(right, left);
    }

    for(int i = 0; i < n; ++i){
        if(flag[i] == 0){
            r = i;
            break;
        }
    }

    Node *root = new Node(r, NULL, NULL);
    findChildren(root);

    queue<Node*> q;
    q.push(root);

    while(!q.empty()){
        levelOrder.push_back(q.front() -> key);
        if(q.front() -> left != NULL) q.push(q.front() -> left);
        if(q.front() -> right != NULL) q.push(q.front() -> right);

        q.pop();
    }

    for(int i = 0; i < n; ++i){
        if(i == n - 1) printf("%d\n", levelOrder[i]);
        else printf("%d ", levelOrder[i]);
    }

    midTraverse(root);
    for(int i = 0; i < n; ++i){
        if(i == n - 1) printf("%d\n", midOrder[i]);
        else printf("%d ", midOrder[i]);
    }

    return 0;
}