ABC 14 D - 閉路 ( Doubling LCA )

D: 閉路 - AtCoder Beginner Contest 014 | AtCoder
A small practice for doubling LCA.

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1e5 + 5;
const int MAXQ = 1e5 + 5;
const int LGN = 20; // 2 ^ 19 > 1e5

int n;
vector<int> G[ MAXN ]; // connects to what
int q;
int a[ MAXN ], b[ MAXN ];

int dpt[ MAXN ]; // depth
int par[ LGN ][ MAXN ]; // par[ i ][ j ]: 2 ^ i ancestor of node j 

void dfs(int u, int fa, int d){
    dpt[ u ] = d;
    par[ 0 ][ u ] = fa;
    for(int v: G[ u ]){
        if( v == fa ) continue;
        dfs( v, u, d + 1 );
    }
}

void preDouble(){ // if par[ i ][ j ] = 0, no such ancestor exists
    for(int i = 1; i < LGN; ++i)
        for(int j = 1; j <= n; ++j){
            int ha = par[ i - 1 ][ j ];
            par[ i ][ j ] = par[ i - 1 ][ ha ];
        }
}

int lca(int u, int v){
    if( dpt[ u ] < dpt[ v ] ) swap<int>( u, v );
    int diff = dpt[ u ] - dpt[ v ];
    for(int k = 0; k < LGN; ++k){
        if( diff & 1 ) u = par[ k ][ u ];
        diff >>= 1;
    }
    if( u == v ) return u;
    for(int k = LGN - 1; k >= 0; --k)
        if( par[ k ][ u ] != par[ k ][ v ] ){
            u = par[ k ][ u ];
            v = par[ k ][ v ];
        }
    assert( par[ 0 ][ u ] == par[ 0 ][ v ] );
    return par[ 0 ][ u ];
}

int dis(int u, int v){ // make sure u or v is the other's root
    return abs( dpt[ u ] - dpt[ v ] );
}

void solve(){
    dfs( 1, 0, 0 ); // let 1 be the root
    preDouble();
    for(int i = 0; i < q; ++i){
        int u = a[ i ], v = b[ i ];
        int ans = dis( u, lca( u, v ) ) + dis( v, lca( u, v ) ) + 1;
        printf("%d\n", ans);
    }
}

int main(){
    scanf("%d", &n);
    for(int i = 0; i < n - 1; ++i){
        int x, y; scanf("%d%d", &x, &y);
        G[ x ].push_back( y );
        G[ y ].push_back( x );
    }
    scanf("%d", &q);
    for(int i = 0; i < q; ++i)
        scanf("%d%d", &a[ i ], &b[ i ]);
    solve();
    return 0;
}