7-9 二叉搜索树的2层结点统计 (25 分)

#include 
#include 
using namespace std;
const int N = 1010, INF = 0x3f3f3f3f;
int l[N], r[N], val[N], n;
int cnt[N], max_depth;

void dfs(int root, int u, int depth) {
    if(val[u] <= val[root]) {
        if(l[root] == INF) {
            l[root] = u;
            cnt[depth+1]++;
            max_depth = max(max_depth, depth+1);
        } else {
            dfs(l[root], u, depth+1);
        }
    } else {
        if(r[root] == INF) {
            r[root] = u;
            cnt[depth+1]++;
            max_depth = max(max_depth, depth+1);
        } else {
            dfs(r[root], u, depth + 1);
        }
    }
}
int main() {
    cin >> n;
    memset(l, 0x3f, sizeof l), memset(r, 0x3f, sizeof r);
    if(n==1) {
        cout << 1 << endl;
        return 0;
    }
    for(int i = 0; i < n; i++) {
        int x; cin >> x;
        if(!i) {
            val[0] = x;
        } else {
            val[i] = x;
            dfs(0, i, 0);
        }
    }
    cout << cnt[max_depth] + cnt[max_depth - 1] << endl;
    return 0;
}