1115. Counting Nodes in a BST (30)

#include <iostream>
#include <cstdio>
#include <functional>
using namespace std;
struct TreeNode
{
int val;
TreeNode *left,*right;
TreeNode(int x):val(x),left(NULL),right(NULL){}
};
void insert(int val,TreeNode *&r)
{
if(r==NULL)
{
TreeNode *p=new TreeNode(val);
r=p;
return;
}
else if(r->val<val)
insert(val,r->right);
else
insert(val,r->left);
}
void solve(TreeNode *r,int lev,const int lowest,int &n1,int &n2)
{
if(r==NULL) return;
if(lev==lowest) n1++;
if(lev==lowest-1) n2++;
solve(r->left,lev+1,lowest,n1,n2);
solve(r->right,lev+1,lowest,n1,n2);
}
int main(int argc, char const *argv[])
{
int n,val;
TreeNode *r=NULL;
cin>>n;
while(n--)
{
cin>>val;
insert(val,r);
}
function<int(TreeNode*)> h=[&h](TreeNode *r){return r==NULL?0:max(h(r->left),h(r->right))+1;};
int lowest=h(r);
int n1=0,n2=0;
solve(r,1,lowest,n1,n2);
printf("%d + %d = %d",n1,n2,n1+n2);
return 0;
}

#include <stdio.h>
#include <malloc.h>
int n1,n2;
struct Node
{
struct Node *left,*right;
int data;
}*T=NULL;
void CreatBST(struct Node **T,int v)
{
if(*T)
{
if(v>(*T)->data)CreatBST(&(*T)->right,v);
else CreatBST(&(*T)->left,v);
}
else
{
(*T)=(struct Node*)malloc(sizeof(struct Node));
(*T)->left=(*T)->right=NULL;
(*T)->data=v;
}
}
void GetHeight(struct Node *T,int level,int *height)
{
if(T)
{
*height=(level>*height)?level:*height;
GetHeight(T->left,level+1,height);
GetHeight(T->right,level+1,height);
}
}
void GetN1AndN2(struct Node *T,int height)
{
struct Node *queue[10000],*s;
int front=0,rear=0,first=0,last=1,h=0,cnt=1;
queue[rear++]=T;
while(front!=rear)
{
s=queue[front++];
++first;
if(s->left){queue[rear++]=s->left;++cnt;}
if(s->right){queue[rear++]=s->right;++cnt;}
if(first==last)
{
++h;last=cnt;
if(h==height-2)n2=cnt-first;
if(h==height-1)n1=cnt-first;
}
}
}
int main()
{
int n,height=0;
scanf("%d",&n);
for(int i=0;i<n;++i)
{
int v;
scanf("%d",&v);
CreatBST(&T,v);
}
GetHeight(T,1,&height);
if(height==1)printf("1 + 0 = 1");
else if(height==2)printf("%d + 1 = %d",n-1,n);
else
{
GetN1AndN2(T,height);
printf("%d + %d = %d\n",n1,n2,n1+n2);
}
return 0;
}