1151 LCA in a Binary Tree

link

 

 

#include <cstdio>
#include <cmath>
#include <vector>
#include <algorithm>
#include <climits>
#include <unordered_map>
#include <cstdio>
#include <iostream>

# define LL long long
using namespace std;

struct TreeNode{
    int val;
    TreeNode* left;
    TreeNode* right;
    TreeNode(int n):val{n},left{NULL},right{NULL}{}
};

int M,N;
int p, q;
TreeNode* buildTree( vector<int> &inorder,  vector<int> &preorder,unordered_map<int,int> &mp, int inleft, int inright, int preleft, int preright){
    if(inleft>inright || preleft>preright) return NULL;

    int val=preorder[preleft];
    int pos=mp[val];
    int leftnums=pos-inleft;
    TreeNode* res=new TreeNode(val);
    res->left=buildTree(inorder,preorder,mp,inleft,pos-1,preleft+1,preleft+1+leftnums-1);
    res->right=buildTree(inorder,preorder,mp,pos+1,inright,preleft+leftnums+1,preright);
    return res;
}

TreeNode* lca(TreeNode* root){
    if(root==NULL) return NULL;
    if(root->val==p || root->val==q) return root;
    TreeNode* left=lca(root->left);
    TreeNode* right=lca(root->right);
    if(left!=NULL && right!=NULL) return root;
    if(left!=NULL) return left;
    return right;
}

int main(){
    scanf("%d %d", &M, &N);
    vector<int> inorder(N);
    vector<int> preorder(N);
    unordered_map<int,int> mp; // inorder value to index
    for(int i=0;i<N;i++){
        scanf("%d", &inorder[i]);
        mp[inorder[i]]=i;
    }
    for(int i=0;i<N;i++){
        scanf("%d", &preorder[i]);
    }
    TreeNode* root=buildTree(inorder,preorder,mp,0,N-1,0,N-1);
    for(int i=0;i<M;++i){
        scanf("%d %d", &p, &q);
        if(mp.find(p)==mp.end() && mp.find(q)==mp.end()){
            printf("ERROR: %d and %d are not found.\n", p, q);
        }else if(mp.find(p)==mp.end()){
            printf("ERROR: %d is not found.\n", p);
        }else if(mp.find(q)==mp.end()){
            printf("ERROR: %d is not found.\n", q);
        }else{
            TreeNode* res=lca(root);
            if(res->val!=p && res->val!=q){
                printf("LCA of %d and %d is %d.\n", p, q, res->val);
            }else{
                if(res->val==p){
                    printf("%d is an ancestor of %d.\n", p,q);
                }else{
                    printf("%d is an ancestor of %d.\n", q,p);
                }
            }
        }
    }

    return 0;
}