백준 13896 - Sky Tax

문제링크 : https://www.acmicpc.net/problem/13896

R이 수도로 주어진다.
1번 쿼리가 수도를 U로 옮기는 것
2번 쿼리가 도시 U의 Tax를 계산하는 것

2번 쿼리에 대한 출력을 해주면 되는데

lca와 조건분기로 풀이를 했따

R이 U와 같으면 그냥 모든도시니깐 N출력

lca(R,U) 가 U와 같으면
U의 밑의 것들을(바뀐 트리에서의 상황을 말한다) 출력해주면되는데
매 질의마다 set_tree를 해줄 수 는 없으니간
R - U 경로에서 U의 바로 밑 자식의 자식의 개수들을 세주면 된다.
이를 위하여 childSum 배열을 활용했으며 dfs로 초기화 시켜줬다.

이후
R - U 사이 노드수(k)를 구해서(두 노드간의 level의 차 )
R 의 k-1 번째 조상을 구하는데 binary lifting 기법?.. 을 사용했다
(아직도 sparse table이 익숙하지가 않다..)

lca(R,U) 가 U하고 다른 상황이라면
그냥 U의 자식들 개수 출력하면된다

#include<iostream>
#include<algorithm>
#include<vector>
#include<cmath>

using namespace std;
int T;
int N;int Q;int R;
vector<int> v[100001];
int p[100001][18];int level[100001];
int max_level;
int childsum[100001];
int dfs(int n){
    int cnt = 1;
    for(int next : v[n]){
        if(level[next]!=0) continue;
        level[next]=level[n]+1;
        cnt += dfs(next);
    }
    return childsum[n]=cnt;
}
void set_tree(int node,int pnode,int lv){
    p[node][0]=pnode;
    level[node]=lv;
    for(int i=1;i<=max_level;i++){
        p[node][i]=p[p[node][i-1]][i-1];
    }
    for(auto a : v[node]){
        int child = a;
        if(child == pnode) continue;
        set_tree(child,node,lv+1);
    }
}
int lca(int a,int b){
    if(a==1||b==1) return 1;
    int target = a; int compare = b;
    if(level[target]<level[compare]){
        swap(target,compare);
    }
    if(level[target]!=level[compare]){
        for(int i=max_level;i>=0;i--){
            if(level[p[target][i]]>=level[compare]){
                target=p[target][i];
            }
        }
    }
    int ret = target;
    if(target!=compare){
        for(int i=max_level;i>=0;i--){
            if(p[target][i]!=p[compare][i]){
                target=p[target][i];
                compare=p[compare][i];
            }
            ret = p[target][i];
        }
    }
    return ret;
}
int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    cin>>T;int cnt = 0;
    max_level=(int)floor(log2(100001));
    while(T>0){
        T--;cnt++;
        cin>>N>>Q>>R;

        for(int i=0;i<100001;i++){
            fill_n(p[i],18,0);
        }
        for(int i=0;i<100001;i++){
            v[i].clear();
        }
        fill_n(level,100001,0);
        fill_n(childsum,100001,0);

        for(int i=0;i<N-1;i++){
            int a,b;
            cin>>a>>b;
            v[a].push_back(b);
            v[b].push_back(a);
        }
        level[1]=1;
        dfs(1);
        set_tree(1,0,1);
        cout<<"Case #"<<cnt<<":\n";
        for(int i=0;i<Q;i++){
            int S,U;
            cin>>S>>U;
            int ans;

            int node = lca(R,U);

            if(S==0){
                R=U;
            }else{
                if(R==U){
                    ans=N;
                }else if(node == U){
                    int diff = level[R]-level[U]-1;
                    int t = R;
                    for(int i=0;i<=max_level;i++){
                        if(diff&1<<i){
                            t=p[t][i];
                            diff-=1<<i;
                        }
                    }
                    ans = N-childsum[t];
                }else{
                    ans = childsum[U];
                }
                cout<<ans<<'\n';
            }
        }
    }
}