PAT---A1115 Counting Nodes in a BST (30分)

题意

给出n个整数依次插入BST，求最后两层的节点数并按照要求输出。‘

思路

构造BST，然后遍历一遍得到各层的节点数。

Sample Input:

9
25 30 42 16 20 20 35 -5 28

Sample Output:

2 + 4 = 6

#include "bits/stdc++.h"
using namespace std;

struct Node{
    int v;
    Node *l{nullptr};
    Node *r{nullptr};
};
Node* insert(Node* node, int v){
    if (node == nullptr) return new Node{v};
    else if (v <= node->v) node->l =  insert(node->l,v);
    else node->r = insert(node->r,v);
    return node;
}

map<int,int,greater<>> layer;
void dfs(Node* node,int depth){
    if (node == nullptr) return;
    layer[depth] ++;
    dfs(node->l,depth+1); dfs(node->r,depth+1);
}
int main(){
//     freopen("input.txt","r",stdin);
    int n; cin >> n;
    Node *root = nullptr;
    for (int i = 0; i < n; ++i) {
        int v; scanf("%d",&v); root = insert(root,v);
    }
    dfs(root,1);
    auto it = layer.begin();
    int n1 = (*it).second; it++;
    int n2 = (*it).second;
    printf("%d + %d = %d\n",n1,n2,n1+n2);
}