PAT甲级 1115 二叉搜索树最后两层结点数量

原题链接

二叉搜索树 (BST) 递归定义为具有以下属性的二叉树：

若它的左子树不空，则左子树上所有结点的值均小于或等于它的根结点的值
若它的右子树不空，则右子树上所有结点的值均大于它的根结点的值
它的左、右子树也分别为二叉搜索树
将一系列数字按顺序插入到一个空的二叉搜索树中，然后，请你计算结果树的最低两层的结点个数。

输入格式
第一行包含整数 N，表示插入数字序列包含的数字个数。

第二行包含 N 个整数，表示插入数字序列。

输出格式
以如下格式，在一行中，输出结果树的最后两层的结点数：

n1 + n2 = n
n1 是最底层结点数量，n2 是倒数第二层结点数量，n 是它们的和。

数据范围
1≤N≤1000,
−1000≤ 插入数字 ≤1000。

输入样例：
9
25 30 42 16 20 20 35 -5 28
输出样例：
2 + 4 = 6

我的解法：

#include <bits/stdc++.h>
using namespace std;
const int N = 1010;
int l[N], r[N], v[N];
int idx, root;
int max_dep, cnt[N];
void insert(int &u, int w){
    if(!u){
        u = ++idx;
        v[u] = w;
    }
    else if(w <= v[u]) insert(l[u], w);
    else if(w > v[u]) insert(r[u], w);
}
void dfs(int u, int depth){
    if(!u) return;
    cnt[depth] ++;
    max_dep = max(depth, max_dep);
    dfs(l[u], depth + 1);
    dfs(r[u], depth + 1);
}
int main(){
    int n;
    cin >> n;
    for(int i = 0; i < n; i ++ ){
        int w;
        cin >> w;
        insert(root, w);
    }
    dfs(root, 0);
    int n1 = cnt[max_dep], n2 = cnt[max_dep - 1];
    printf("%d + %d = %d", n1, n2, n1 + n2);
    return 0;
}

收获：

解决此类问题的关键是建立树

灵活运用 & ，双向修改