백준 13511번 : 트리와 쿼리 2 (C/C++)

 

<문제>

백준 13511번

주어진 쿼리에 대해 답하는 문제이다.

 

<참고한 문제>

백준 1761번: 정점들의 거리

백준 11438 :  LCA 2

 

<풀이>

아래 배열을 정의한다.

• depth[x] : 루트로부터 정점 x까지의 거리
• lev[x] : 루트를 기준으로 정점x의 레벨
• pa[x][i] : 정점 x의 2^i번째 부모

우선 dfs탐색을 통해 depth, lev, pa배열을 채운다.

l = lca(u, v)라고 할때, 쿼리1은 (depth[u]-depth[l])+(depth[v]-depth[l])로 구할 수 있다.

 

쿼리 2는 u부터 시작하여 k번 탐색을 진행하여야하는데 매 쿼리마다 k번의 탐색을 하면 시간초과가 날 것이다.

효율적인 풀이를 위해 경우를 나눈다.

 

    1.  k < lev[u]-lev[l]+1 인 경우 :  u의 k-1번째 부모를 구한다.

 

     2. k >= lev[u]-lev[l]+1 인경우 : v의 (lev[u]-lev[l])+(lev[v]-lev[l])-k+1 번째 부모를 구한다.

climb(u, t) 를 u의 t번째 부모를 구하는 함수로 정의한다.

 

<코드>

#include<bits/stdc++.h>
#define fio()                     \
    ios_base::sync_with_stdio(0); \
    cin.tie(0)
using namespace std;
 
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef tuple<int, int, int> tpi;
typedef tuple<ll, ll, ll> tpl;
typedef pair<double, ll> pdl;
 
const int INF = 0x3f3f3f3f;
const ll LINF = 0x3f3f3f3f3f3f3f3f;
const int dx[] = { 0, 1, 0, -1 };
const int dy[] = { 1, 0, -1, 0 };
 
vector<pii> g[101010];
int lev[101010], pa[101010][17];
ll depth[101010];
 
void dfs(int x, int d){
    lev[x] = d;
 
    for(auto [nx, dist] : g[x]){
        if(lev[nx]) continue;
        depth[nx] = depth[x] + (ll) dist;
        pa[nx][0] = x;
        dfs(nx, d+1);
 
    }
    return;
}
 
int lca(int x, int y){
    if(lev[x] < lev[y]) swap(x, y);
 
    int diff = lev[x] - lev[y];
    int j = 0;
    while(diff){
        if(diff&1) x = pa[x][j];
        j++;
        diff >>= 1;
    }
 
    for(int i = 16; i >= 0; i--){
        if(pa[x][i] != pa[y][i]){
            x = pa[x][i];
            y = pa[y][i];
        }
    }
 
    while(x != y){
        x = pa[x][0];
        y = pa[y][0];
    }
 
    return x;
}
 
int climb(int x, int t){
    int j = 0;
    while(t){
        if(t&1) x = pa[x][j];
        t >>= 1;
        j++;
    }
    return x;
}
 
int main(){
    int n; scanf("%d", &n);
    for(int i = 0; i < n-1; i++){
        int u, v, w; scanf("%d %d %d", &u, &v, &w);
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }
    dfs(1, 1);
    for(int i = 1; i <= 16; i++){
        for(int j = 1; j <= n; j++){
            pa[j][i] = pa[pa[j][i-1]][i-1];
        }
    }
    int m; scanf("%d", &m);
    for(int i = 0; i < m; i++){
        int q, u, v; scanf("%d %d %d", &q, &u, &v);
        int l = lca(u, v);
        if(q == 1){
            printf("%lld\n", (depth[u]-depth[l])+(depth[v]-depth[l]));
        }
        else{
            int ans, k; scanf("%d", &k);
            if(k < lev[u]-lev[l]+1){
                ans = climb(u, k-1);
            }
            else{
                ans = climb(v, (lev[u]-lev[l])+(lev[v]-lev[l])+1-k);
            }
            printf("%d\n", ans);
        }
    }
}

<회고>

lca2 알고리즘을 활용해볼 수 있었다.