1151 LCA in a Binary Tree (30 分)

考前挣扎一下

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

const int maxn=10010;
map<int,int> mp,index_k,height;
int m,n,pre[maxn],in[maxn];

void build(int level,int father,int prel,int prer,int inl,int inr){
    if(prel>prer){
        return;
    }
    int root=pre[prel];
    mp[root]=father;
    height[root]=level;
    int i;
    for(i=inl;i<=inr;i++){
        if(in[i]==root){
            break;
        }
    }
    int cnt=i-inl;
    build(level+1,index_k[root],prel+1,prel+cnt,inl,i-1);
    build(level+1,index_k[root],prel+cnt+1,prer,i+1,inr);
}

int main(){
    cin>>m>>n;
    for(int i=0;i<n;i++){
        scanf("%d",&in[i]);
        //a[i]=in[i];
        index_k[in[i]]=i;
    }
    for(int i=0;i<n;i++){
        scanf("%d",&pre[i]);
    }
    build(0,-1,0,n-1,0,n-1);
    int n1,n2;
    for(int i=0;i<m;i++){
        cin>>n1>>n2;
        if(index_k.find(n1)==index_k.end()&&index_k.find(n2)==index_k.end()){
            cout<<"ERROR: "<<n1<<" and "<<n2<<" are not found."<<endl;
        }
        else if(index_k.find(n1)==index_k.end()){
            cout<<"ERROR: "<<n1<<" is not found."<<endl;
        }
        else if(index_k.find(n2)==index_k.end()){
            cout<<"ERROR: "<<n2<<" is not found."<<endl;
        }
        else{
            int big=n1,small=n2;
            if(height[n1]<height[n2]){
                small=n1;
                big=n2;
            }
            int bb=big,ss=small;
            int diff=height[big]-height[small];
            for(int j=0;j<diff;j++){
                big=in[mp[big]];
            }
            int times=0;
            while(1){
                if(big==small){
                    if(times==0){
                        cout<<ss<<" is an ancestor of "<<bb<<"."<<endl;
                    }
                    else{
                        cout<<"LCA of "<<n1<<" and "<<n2<<" is "<<big<<"."<<endl;
                    }
                    break;
                }
                big=in[mp[big]];
                small=in[mp[small]];
                times++;
            }
        }
    }
    return 0;
}