[코드트리] 노드의 공통 조상 2 C++

문제

https://www.codetree.ai/missions/9/problems/common-ancestor-of-node-2?&utm_source=clipboard&utm_medium=text

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

코드

#include <iostream>
#include <vector>
using namespace std;

const int MAX_N = 50'001;
const int MAX_H = 16; // 2^16=65,536 이므로 50,000개의 노드를 커버할 수 있다

int n, q;
vector<int> edges[MAX_N];
bool visited[MAX_N]={false, };

int depth[MAX_N];
int parent[MAX_H][MAX_N];

// DFS를 통해 parent 초기 조건을 채움
void DFS(int x){
    for(int y : edges[x]){
        if(visited[y]) continue;

        visited[y]=true;
        parent[0][y]=x; // y에서 1(2^0)번 올라간 노드는 x(=x가 부모라는 의미)
        depth[y]=depth[x]+1;
        DFS(y);
    }
}

int LCA(int a, int b){
    if(depth[a] < depth[b]) swap(a, b);

    // 1. 두 노드의 깊이를 같게 만들기
    for(int h=MAX_H; h>=0; --h){
        if(depth[a]-depth[b] >= (1<<h)) a=parent[h][a];
    }
    
    if(a == b) return a;

    // 2. 두 노드의 공통 조상 찾기(두 노드가 일치하기 직전까지 올라감)
    for(int h=MAX_H; h>=0; --h){
        if(parent[h][a] != parent[h][b]){
            a=parent[h][a];
            b=parent[h][b];
        }
    }

    return parent[0][a];
}

int main() {
    cin >> n;
    for(int i=1; i<n; ++i){
        int a, b;
        cin >> a >> b;
        edges[a].push_back(b);
        edges[b].push_back(a);
    }
    
    // parent 초기조건 채우기(루트는 1번)
    depth[1]=1;
    visited[1]=true;
    DFS(1);

    // 점화식을 이용하여 parent 값 구하기
    for(int h = 1; h < MAX_H; h++) {
        for(int i = 1; i <= n; i++) {
            parent[h][i] = parent[h-1][parent[h-1][i]];
        }
    }

    // q번 연산
    cin >> q;
    while(q--){
        int a, b;
        cin >> a >> b;
        cout << LCA(a, b)  << '\n';
    }

    return 0;
}

 

개념

LCA(Lowest Common Ancestor)

노드 a와 b의 조상 중 깊이가 가장 깊은 조상(=a, b가 리프에서 올라갔을 때 가장 먼저 만나게 되는 공통인 조상)

 

노드의 인접한 공통 조상 문제에서 LCA를 찾기 위해 2개의 과정을 거쳤다.

1. a, b 노드의 깊이가 다를 경우 깊이가 적은 쪽에 깊이를 맞춘다.
2. 두 노드가 같아질 때까지 노드를 한 칸 씩 올라간다

// depth가 적은 노드의 깊이에 맞춤
while(depth[a]!=depth[b]){
    if(depth[a]>depth[b]) a=parent[a];
    else b=parent[b];
}

// 한 칸씩 같이 올라가면서 두 노드가 같아질 때까지 반복
while(a!=b){
    a=parent[a];
    b=parent[b];
}

 

이 방식은 최악의 경우 모든 노드(n개)를 탐색해야 할 수 있다. 따라서 $O(n)$의 시간복잡도를 가진다.

 

 

 

Sparse Table(그냥 특별한 배열이라고 생각하면 됨)를 사용하면 $O(logN)$에 LCA을 찾을 수 있다.

1. 2차원 배열 parent[h][i]  채우기

parent[h][i] = i번 노드에서 2^h번 부모를 따라 위로 올라갔을 때의 노드 번호를 관리

 

1-1. [초기 조건] parent[0][i]= i번 노드에서 $2^0(=1)$번 부모를 따라 위로 올라 갔을 때의 노드 번호. 즉, i번 노드의 부모 노드 번호.

1-2. [점화식] parent[h][i]=parent[h-1][parent[h-1][i]]

parent 배열을 채우는 데에는 $O(NlogN)$이 걸린다.

$2^h$번 위로 올라갔을 때의 위치를 구해야 하는데, h의 값은 최대 logN

logN * N 크기의 테이블에 값을 채워야 하므로 시간복잡도는 $O(NlogN)$

 

2. parent[h][i]를 이용하여 LCA 찾기

2-1. 두 노드의 깊이 맞추기

두 노드의 깊이 차이를 구해서, 깊이가 작은쪽에 맞춘다.

a의 깊이가 5이고, b의 깊이가 18이라면 둘의 깊이 차이는 13이 된다.

 

이때 방법이 2가지가 있다.

첫 번째 방법은 위에서 했던 것처럼 1칸씩 올리는 방법이다.

b=parent[b];

일직선으로 이어진 트리에서 a가 루트노드이고, b가 리프 노드일 경우(최악의 경우), n번의 연산을 해야한다. $O(n)$

(이 방법을 쓰면 parent 배열을 만든 의미가 없어진다...)

 

두 번째 방법은 십진수를 이진수로 변환하는 방법이다.

13을 이진수로 바꾸면 1101이다.

$1 × 2^3 + 1 × 2^2 + 0 × 2^1 + 1 × 2^0 = 8 + 4 + 0 + 1 = 13$

1101은 8+4+1이다. 2의 가장 큰 거듭제곱(여기서는 8)부터 바로 계산한다.

parent의 정의(parent[h][i]는 i번 노드에서 $2^h$만큼 올라간 노드)에 따라 13에서 8($2^3$)만큼 올라간 노드는 parent[3][b]가 될 것이다. b 노드에서 $2^3=8$만큼 올라간 노드이기 때문.

 

13에서 8을 빼면 5가 남는다. 5의 이진수는 101이고,

$1 × 2^2 + 0 × 2^1 + 1 × 2^0 = 8 + 4 + 0 + 1 = 13$

4+1이므로, 2의 가장 큰 거듭제곱은 4현재 b에서 4만큼 올라간 노드는 parent[2][b]: b노드에서 $2^2$만큼 올라간 노드

 

5에서 4를 빼면 1이고, 1의 이진수는 1이다.

2의 가장 큰 거듭제곱은 1($2^0$)

b에서 1만큼 올라간 노드는 parent[0][b]: b노드에서 $2^0$만큼 올라간 노드

 

이제 a와 b의 높이는 같아졌다. 방금 과정을 다시 살펴보자.

1101(=13의 이진수)에서 101(=5의 이진수)가 되었고, 1이 되었다.

'<<' shift연산을 이용하면 되는 것!(2의 제곱 연산에는 shift를 사용할 수 있음을 기억하자!)

8->4->2->1 이런식으로 탐색해야할 값이 제곱으로 줄어들기 때문에 시간복잡도는 $O(logN)$이 된다.

 

2-2. 공통 조상 찾기

공통 조상도 parent배열을 이용해서 찾으면 된다.

h의 최대값부터 점점 내려오면서 최초로 일치하지 않는 높이를 찾는다.(반복문)

그리고 그 높이에서 1(2^0)칸 올라가면 두 노드는 일치하게 된다. a=parent[0][a]