[C/C++] 백준 #1761 정점들의 거리(최소 공통 조상)

트리에서 두 노드간 최소 거리를 구하는 방법은 여러가지가 있습니다.  가장 간단하게 생각할 수 있는 것은 깊이 우선 탐색(DFS) 이용하는 방법이 있겠죠.  하지만, 노드의 수가 많다면, 꽤 오랜 시간이 걸리겠지만, 크게 문제가 되지 않습니다.  노드의 개수가 N개라면,  \(O(N)\) 알고리증으로 풀 수 있으니까요.  물론 재귀 호출의 깊이가 깊어지기 때문에 너비 우선 탐색(BFS)를 이용할 수도 있습니다.

 

하지만, 이번 문제는 단순하게 두 노드간의 거리를 한번 구하는 것이 아니고 M번 구하는 것입니다.  그러면, 알고리즘 효율이 \(O(NM)\)이 됩니다.  그렇기 때문에 보다 효율적인 알고리즘을 찾아야 합니다.  최소 공통 조상 알고리즘은 기본적으로 많은 질의(query)가 있는 경우 사용할 수 있습니다.

 

모든 노드는 자신의 조상에 대한 정보를 저장합니다.  그런데, 이 정보를 단순하게, 나의 K대 조상까지 저장하는 방식이 아닙니다.  최악의 경우 양쪽으로 나누어진 트리의 경우에는 이 방식으로 해도 결국 알고리즘 효율은 \(O(NM)\)을 벗어나지 못 합니다.

 

two splited tree

위와 같은 트리라면, 왼쪽에 있는 노드에서 오른쪽에 있는 노드간에 거쳐야 하는 노드의 개수가 항상 많겠죠.  지금은 9개의 노드가 있지만, 이 노드의 수가 많다면, 결국 많은 시간이 걸릴겁니다.

 

그래서 미리 작업을 해서 조상에 대한 정보를 각 노드별로 놓습니다.  단순하게 K대 조상을 차례대로 나열해도 되겠지만, 트리의 깊이가 깊은 경우에는 그렇게 효율적이지 않습니다.  K를 무지막지하게 크게 한다면, 가능합니다.  예를 들어서 6번 노드의 조상의 [ 4, 2, 1 ] 입니다.  8번 노드는 부모 노드의 조상들에 앞에 부모 노드인 6을 넣어서 [ 6, 4, 2, 1 ]을 만들면 됩니다.  하지만 위의 그림과 같은 경우에는 최대 N개의 조상들을 저장할 공간을 가지고 있어야 합니다.  메모리 효율이 \(O(N^2)\)이 됩니다.  그리고 공통 조상을 찾는 것도 \(O(N)\)이 되므로 바람직하지 않습니다.

 

그래서 사용하는 방법이 \(2^k\)번째에 해당하는 조상들만 저장하는 것입니다.  그러면 메모리 효율은 \(O(N \log N)\)이 됩니다.  또한 검색을 하는 속도도 빨라지게 됩니다.

 

조상을 만드는 작업은 너비 우선탐색을 사용합니다.  그리고 루트(root)에서 얼마나 멀리 떨어져 있는지도 저장을 해야 합니다.  루트는 조상이 없기 때문에 비어있는 조상 리스트를 가집니다.  \(a_{i, j}\) 를 \(i\)번째 노드의 \(2^j\)번째 조상 노드라고 한다면, 다음과 같이 식을 세울 수 있습니다.  \(p_{i}\)는 \(i\)번째 노드의 부모입니다.

 

\[ a_1 = [ ] \]

\[ a_{i, 0} = p_{i} \]

\[ a_{i, j} = a_{a_{i, j-1}, j-1} \]

 

이 식에 따라서 위의 그래프를 해석하면, 루트 노드인 1번 노드는 [ ] 조상 리스트를 가집니다.   2번 노드는 [ 1 ] 을 조상으로 가지게 되죠.  4번 노드는 [ 2, 1 ]을 가지게 됩니다.  6번 노드는 [ 4, 2 ] 을 가집니다.  8번 노드는 [ 6, 4, 1 ]을 가지게 됩니다.

 

공통 조상을 찾기 위해서는 두 노드의 트리의 깊이를 서로 맞추어야 합니다.  위의 식을 이용하면, 이진탐색과 동일하게 접근해서 원하는 트리의 깊이의 조상으로 접근할 수 있습니다.

 

같은 깊이가 되었으면, 조상들 리스트가 처음으로 달라지는 지점을 선택해서 그 지점으로 두개의 노드를 이동합니다.  그리고, 또 조상들 리스트를 검사해보는 식으로 반복하면 해답을 찾을 수가 있습니다.

 

제가 작성한 소스입니다.  소스는 참고용으로 봐주세요.

//----------------------------------------------------------
//    baekjoon #1761 - Distance between two nodes
//        - by Aubrey Choi
//        - created at 2020-01-22
//----------------------------------------------------------
#include <stdio.h>
#include <memory.h>
#include <vector>
#include <queue>

static int dp[100001][18], w[100001];
static std::vector<int> adj[100001];
int d[100001], n;
void bfs(int r)
{
    int i, s;
    std::queue<int> q;
    q.push(r); d[r]=1; memset(dp[r],-1,sizeof(int)*18);
    while(!q.empty())
    {
        s=q.front(); q.pop();
        for(auto k:adj[s])
        {
            if(d[k/20000]) continue;
            w[k/20000]=w[s]+k%20000; k/=20000;
            d[k]=d[s]+1;
            for(i=0,dp[k][i]=s;dp[k][i]!=-1;i++) dp[k][i+1]=dp[dp[k][i]][i];
            memset(dp[k]+i,-1,sizeof(int)*(18-i));
            q.push(k);
        }
    }
}
int main()
{
    int i, a, b, c, m, t, r;
    scanf("%d",&n);
    for(i=1;i<n;i++)
    {
        scanf("%d%d%d",&a,&b,&c);
        adj[a].push_back(b*20000+c);
        adj[b].push_back(a*20000+c);
    }
    bfs(1);
    scanf("%d",&m);
    while(m--)
    {
        scanf("%d%d",&a,&b);
        r = w[a]+w[b];
        if(d[a]<d[b]) { i=a; a=b; b=i; }
        for(t=d[a]-d[b],i=0;t;t>>=1,i++) a=(t&1)?dp[a][i]:a;
        while(a!=b)
        {
            for(i=0;dp[a][i]!=dp[b][i];i++);
            if(i<2) { a=dp[a][i]; break; }
            a=dp[a][i-1], b=dp[b][i-1];
        }
        printf("%d\n",r-w[a]*2);
    }
}