least common ancestor (doubling)

aoj_grl_5_c.cpp

#include <bits/stdc++.h>
#define REP(i, n) for(int i=0; i < (int)(n); i++)
using namespace std;

typedef long long ll; // -9,223,372,036,854,775,808 から 9,223,372,036,854,775,807

template <class T> using vec = std::vector<T>;
template <class T> using vec2 = vec<vec<T>>;
template <class T> using vec3 = vec<vec<vec<T>>>;

// #define DEBUG

int main(){
    int n;
    cin >> n;
    vec<int> up(n, -1);
    for(int i=0; i < n; i++){
        int Nchild;
        cin >> Nchild;
        for(int j=0; j < Nchild; j++){
            int c;
            cin >> c;
            c;
            up.at(c) = i;
        }
    }

    int q;
    cin >> q;
    vec<complex<int>> queries;
    for(int i=0; i < q; i++){
        int u, v;
        cin >> u >> v;
        queries.push_back(complex<int>(u, v));
    }

    // preprocess
    int logN = int(log(n)/log(2)) + 1;
    vec2<int> upDouble(logN, vec<int>(n, -1)); 
    // upDouble[i][j] means the position 2^i nodes above j
    for(int j=0; j < n; j++){
        upDouble.at(0).at(j) = up.at(j);
    }

    for(int i=1; i < logN; i++){
        for(int j=0; j < n; j++){
            int posUpByI = upDouble.at(i-1).at(j);
            if(posUpByI == -1){
                continue;
            }
            posUpByI = upDouble.at(i-1).at(posUpByI);
            if(posUpByI == -1){
                continue;
            }
            upDouble.at(i).at(j) = posUpByI;
        }
    }

    #ifdef DEBUG
    printf("*START: up\n");
    for(int i=0; i < upDouble.size(); i++){
        for(int j=0; j < upDouble[0].size(); j++){
            cout << upDouble[i][j] << "\t";
        }
        cout << endl;
    }
    printf("*END: up\n");
    #endif

    auto getIAbove = [&upDouble](int pos, int i){
        int deg = 0;
        while(i > 0){
            if(i & 1){
                pos = upDouble.at(deg).at(pos);
            }
            i >>= 1;
            deg++;
        }
        return pos;
    };

    #ifdef DEBUG
    printf("*START: test\n");
    cout << getIAbove(7, 2) << endl;
    cout << getIAbove(7, 3) << endl;
    printf("*END: test\n");
    #endif

    for(auto query: queries){
        int u, v;
        u = query.real();
        v = query.imag();

        auto getLen = [&](int pos){
            int res = 0;
            for(int i=logN-1; i >= 0; i--){
                if(upDouble.at(i).at(pos) != -1){
                    pos = upDouble.at(i).at(pos);
                    res += pow(2, i);
                }
            }
            return res;
        };
        // get len(u) and len(v)
        int lenU = getLen(u);
        int lenV = getLen(v);

        // find the maximum i such that i nodes below the root is common ancestor.
        // smaller i becomes TRUE.
        ll bb = 0;
        ll uu = n;
        // f(-)は増加。減少なら不等式は>か>=にする。
        auto check = [&](int below){
            return getIAbove(u, lenU - below) == getIAbove(v, lenV - below);
        };
        while(uu - bb > 1){
            ll piv = (uu + bb)/2;
            if(check(piv)){ // 必要なら不等号はとりかえる。とにかく左側がTRUEということにする(右がTRUEなら反転する)
                bb = piv;
            }else{
                uu = piv;
            }
        }
        // bはfがK未満(条件が真)の位置のうち一番右、uはfがK以上(条件が偽)の位置のうち一番左
        cout << getIAbove(u, lenU - bb) << endl;
    }







    return 0;
}