pat1115 Counting Nodes in a BST

题意：输入一颗bst，输出最后两层的节点数，及其节点数的和。

思路：建树，dfs求深度，层序求节点数。

代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <queue>

using namespace std;

const int MAX_N = 1010;
int n;
int a[MAX_N];
int lc[MAX_N], rc[MAX_N];
int root;

int ins(int r, int x) {
    if (r == -1) return x;
    if (a[x] <= a[r]) lc[r] = ins(lc[r], x);
    else rc[r] = ins(rc[r], x);
    return r;
}

int heigh;
int dfs(int rt) {
    if (rt == -1) return -1;
    int l, r; l = r = -1;
    l = dfs(lc[rt]);
    r = dfs(rc[rt]);
    return max(l, r) + 1;
}

int a1, a2;
int h[MAX_N];
void level(int r) {
    memset(h, 0, sizeof(h));
    h[r] = 0;
    queue<int> q; q.push(r);
    while (!q.empty()) {
        int t = q.front(); q.pop();
        if (h[t] == heigh) a1++;
        if (h[t] == heigh-1) a2++;
        if (lc[t] != -1) {h[lc[t]] = h[t]+1; q.push(lc[t]);}
        if (rc[t] != -1) {h[rc[t]] = h[t]+1; q.push(rc[t]);}
    }
}

int main() {
    // freopen("in.txt", "r", stdin);
    // freopen("out.txt", "w", stdout);
    scanf("%d", &n);
    root = -1;
    memset(lc, -1, sizeof(lc));
    memset(rc, -1, sizeof(rc));
    for (int i = 0; i < n; i++) {
        scanf("%d", &a[i]);
        root = ins(root, i);
    }
    heigh = dfs(root);
    a1 = a2 = 0;
    level(root);
    printf("%d + %d = %d\n", a1, a2, a1+a2);
    return 0;
}