(라이님 블로그 대회 알고리즘 따라잡기 30) LCA(Lowest Common Ancestor)

오늘은 LCA를 배워보았다. 그 전에 Sparse Table에 대해 알아두면 좋을 것 같아 Sparse Table에 대해서도 공부해보았다. 라이님 LCA 포스팅을 보다보니 Sparse Table이 나왔는데 전에 본 적이 없어서 이 참에 공부!

 

Sparse Table에 대한 설명은 블로그를 참고하고, 아래는 아주 조금 보기 쉽게 고친 코드이다.

#include <cstdio>
using namespace std;
const int MAX = 500001;
const int MAX_D = 19; // 2^k >= MAX인 최소의 k

int main() {
	// next의 크기가 너무 클 경우 vector<vector<int>>로 정의
    int M, next[MAX][MAX_D];

    // 정점 입력
    scanf("%d", &M);
    for (int i = 1; i <= M; ++i)
        scanf("%d", &next[i][0]);

    // 희소 테이블 생성
    for (int j = 1; j < MAX_D; ++j) {
        for (int i = 1; i <= M; ++i) {
            next[i][j] = next[next[i][j - 1]][j - 1];
        }
    }

    // 쿼리 처리
    int Q;
    scanf("%d", &Q);
    while (Q--) {
        int n, x;
        scanf("%d %d", &n, &x);
        for (int j = 0; j < MAX_D && n > 0; ++j) {
        // n의 j번째 비트가 1인지 확인하는 방법
            if (n & (1 << j)) {
                x = next[x][j];
            }
        }
        // 결과 출력
        printf("%d\n", x);
    }
    return 0;
}

 

 

이제 LCA에 대해 알아보자. 라이님 블로그에 나온 설명으로 LCA의 설명은 갈음한다.

 

#include <cstdio>
#include <vector>
#include <algorithm>
#include <cstring>
using namespace std;
const int MAX = 18; // roundup log(2, 100000)

int N, M;
vector<vector<int>> parent; // parent[i][k]: i의 2^k번째 부모
vector<int> depth; // 정점의 깊이 (루트의 깊이: 0)
vector<vector<int>> adj; // 인접 리스트

// 트리를 DFS로 생성하며 parent[i][0]과 depth 배열을 채움
void makeTreeByDFS(int curr) {
    for (int next : adj[curr]) {
        if (depth[next] == -1) {
            parent[next][0] = curr;
            depth[next] = depth[curr] + 1;
            makeTreeByDFS(next);
        }
    }
}

int main() {
    // 트리 정보 입력
    scanf("%d", &N);
    parent.resize(N, vector<int>(MAX, -1));
    depth.resize(N, -1);
    adj.resize(N);

    for (int i = 0; i < N - 1; ++i) {
        int u, v;
        scanf("%d %d", &u, &v);
        u--; v--;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    // 트리 만들기
    depth[0] = 0;
    makeTreeByDFS(0);

    // parent 배열 채움
    for (int j = 0; j < MAX - 1; ++j) {
        for (int i = 0; i < N; ++i) {
            if (parent[i][j] != -1) {
                parent[i][j + 1] = parent[parent[i][j]][j];
            }
        }
    }

    // 쿼리 입력받기
    scanf("%d", &M);
    for (int i = 0; i < M; ++i) {
        int u, v;
        scanf("%d %d", &u, &v);
        u--; v--;

        // depth[u] >= depth[v]가 되도록 필요에 따라 u, v를 스왑
        if (depth[u] < depth[v]) swap(u, v);
        int diff = depth[u] - depth[v];

        // 깊이 차(diff)를 없애며 u를 이동시킴
        for (int j = 0; j < MAX; ++j) {
            if (diff & (1 << j)) {
                u = parent[u][j];
            }
        }

        // u와 v가 다르면 동시에 일정 높이만큼 위로 이동시킴
        if (u != v) {
            for (int j = MAX - 1; j >= 0; --j) {
                if (parent[u][j] != -1 && parent[u][j] != parent[v][j]) {
                    u = parent[u][j];
                    v = parent[v][j];
                }
            }
            // 마지막엔 두 정점 u, v의 부모가 같으므로 한 번 더 올림
            u = parent[u][0];
        }

        // LCA 출력
        printf("%d\n", u + 1);
    }

    return 0;
}

 

LCA 알고리즘을 이용하면 트리의 두 정점 사이에 거리를 빠르게 구하는 것이 가능하다.