백준[11438] LCA 2

문제 설명

문제 링크

https://www.acmicpc.net/problem/11438

풀이

LCA 기본 문제이다. 선행 문제인 이 문제는 시간 복잡도가 쿼리당 O(N)으로 풀 수 있었지만 이 문제는 시간 초과가 발생한다.  그렇기 때문에 이 문제는 sparse table을 이용하여 조금 더 최적화를 해야 한다. 

sparse table은 2^k =  2^(k-1) + 2^(k-1)인 것을 이용하여 모든 정점에서 2^k 만큼 움직이면 어떤 점으로 이용동하는지 dp와 비슷하게 구해놓는다. 이 전처리는 O(NlogN)에 할 수 있다.

dp[k][i] : i번째 노드에서 2^k 움직이면 도착하는 노드

// sparse table 코드
for (int k = 1; k <= 20; k++) {
        for (int i = 1; i <= n; i++) {
            dp[k][i] = dp[k - 1][dp[k - 1][i]];
        }
    }

쿼리를 입력받았을 때 로직은 O(N)풀이와 동일하지만 구하는 방법이 조금 다르다.

1) 두 노드의 level을 맞춘다

2) 두 노드를 부모가 동일할때 까지 올린다.

두 노드의 level 차가 만약 7이라면 4 -> 2 -> 1씩 움직여 O(logN)의 시간 복잡도에 level을 맞춘다.

if (lev[a] > lev[b]) swap(a, b); // b가 a보다 더 크게 만든다
//20은 log2(100000)보다 큰 수를 임의로 잡았다
for (int i = 20; i >= 0; i--) {
    if (lev[b] - lev[a] >= (1 << i)) b = dp[i][b];
}

level을 동일하게 맞췄다면 아래 그림과 같이 최소 공통 조상보다 더 조상 노드로 간다면 조상노드는 항상 일치한다는 특성을 이용한다.

예시

만약 2^20움직였을 때 두 조상이 다르다면 최소 공통 조상은 2^20보다 위에 있는 것이다. 그렇기 때문에 2^20만큼 움직인다. 하지만 만약 두 조상이 같다면 최소 공통 조상보다 더 조상 노드가 되는 것이기 때문에 움직이지 않는다. 이러한 방법으로 2^20 2^19... 1 씩 움직일 수 있는지 체크하며 움직이면 두 노드는 최소 공통조상 바로 아래 노드에 위치할 것이다. 

    for (int i = 20; i >= 0; i--) {
        if (dp[i][a] != dp[i][b]) {
            a = dp[i][a];
            b = dp[i][b];
        }
    }

전체 코드

#include <cstdio>
#include <vector>

using namespace std;

vector<int> v[100005];
int dp[21][100005], lev[100005];

void dfs(int cur, int prv) {
    dp[0][cur] = prv;
    lev[cur] = lev[prv] + 1;
    for (auto nxt : v[cur]) {
        if (prv != nxt) dfs(nxt, cur);
    }
}

int query(int a, int b) {
    if (lev[a] > lev[b]) swap(a, b);
    for (int i = 20; i >= 0; i--) {
        if (lev[b] - lev[a] >= (1 << i)) b = dp[i][b];
    }

    if (a == b) return a;

    for (int i = 20; i >= 0; i--) {
        if (dp[i][a] != dp[i][b]) {
            a = dp[i][a];
            b = dp[i][b];
        }
    }
    return dp[0][a]; //a의 부모노드를 출력
}

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

    dfs(1, 0);

    for (int k = 1; k <= 20; k++) {
        for (int i = 1; i <= n; i++) {
            dp[k][i] = dp[k - 1][dp[k - 1][i]];
        }
    }

    int m; scanf("%d", &m);
    while (m--) {
        int a, b; scanf("%d %d", &a, &b);
        printf("%d\n", query(a, b));
    }
}