PAT A1115 Counting Nodes in a BST (30 分) 二叉搜索树

    题目大意：给出N个数，依次插入到二叉搜索树中，求出最后两层的节点数目和。

    先按照二叉搜索树进行插入，然后层序遍历得到最大深度和每层的节点数，最后输出即可。

AC代码：

#include <vector>
#include <cstdio>
#include <queue>

using namespace std;

struct Node
{
    int data;
    Node* left;
    Node* right;
    int level;
    Node(int data):data(data), left(NULL), right(NULL){};
};

void insert(Node*& node, int data)
{
    if(node == NULL)
    {
        node = new Node(data);
        return;
    }
    if(data > node->data) insert(node->right, data);
    else insert(node->left, data);
}

void levelOrder(Node* node, vector<int> &cnt, int &maxLevel)
{
    queue<Node*> q;
    node->level = 0;
    q.push(node);
    while(!q.empty())
    {
        Node* now = q.front();
        q.pop();
        if(now->level > maxLevel) maxLevel = now->level;
        cnt[now->level]++;
        if(now->left)
        {
            now->left->level = now->level + 1;
            q.push(now->left);
        }
        if(now->right)
        {
            now->right->level = now->level + 1;
            q.push(now->right);
        }
    }
}

int main()
{
    int N;
    scanf("%d", &N);
    vector<int> v(N);
    Node *root = NULL;
    for (int i = 0; i < N; ++i)
    {
        scanf("%d", &v[i]);
        insert(root, v[i]);
    }
    vector<int> levelCnt(N, 0);
    int maxLevel = 0;
    levelOrder(root, levelCnt, maxLevel);
    printf("%d + %d = %d", levelCnt[maxLevel], levelCnt[maxLevel - 1], levelCnt[maxLevel] + levelCnt[maxLevel - 1]);
    return 0;
}