트리의 지름

Table of Contents

개요

방법 1 - 동적계획법

방법 1 코드

방법 2 - 탐욕법

방법 2 코드

방법 2 의 정당성

방법 2 로 far 배열 찾기

문제

1. 개요

 트리의 지름이란, 트리 내에서 가장 먼 두 정점 사이의 거리를 뜻합니다. 이 글에서는 트리의 지름을 구하는 2가지 방법에 대해 소개합니다. 참고로 이 두 가지 방법은 모두 시간 복잡도를 늘리지 않는 선에서 'far [ v ] = 정점 v와 v에서 가장 먼 정점 사이의 거리'를 모든 정점에 대해 구할 수 있습니다.

 far 배열을 구하면 지름 = MAX(far [ v ])로 쉽게 구할 수 있으며, 지름만 구할 때 보다 더 다양한 상황에 쓸 수 있게 됩니다. 때문에 이 문서에서는 far [ v ] 을 구하는 방법까지 소개하도록 하겠습니다.

 방법 1은 동적 계획법 패러다임에 가까운 방법이고, 방법 2는 그리디 패러다임에 가까운 방법입니다. 방법 2가 비교적 더 간단하므로, 하나만 알고 싶다면 방법 2로 넘어가는 걸 추천합니다. 시간복잡도는 둘 다 선형시간, O(노드의 개수) 입니다.

2. 방법 1 - 동적계획법

 far [ v ] = 정점 v와 v에서 가장 먼 정점 사이의 거리

 위와 같은 배열을 구하면 트리의 지름은 MAX(far[v]) 로 쉽게 구할 수 있습니다. 이 아이디어를 기반으로 시작하면 아래와 같이 DP스러운 생각을 전개할 수 있습니다.

 일단 위의 트리에서 far [ 2 ] 에 대해 생각해 보겠습니다. 2를 기준으로 정점은 두 가지로 분류할 수 있지요. 2의 서브 트리에 있는 노드들과 없는 노드들입니다.

dist [ v ] = 정점 v와 v에서 가장 먼 서브트리 내 정점 사이의 거리

parent_far [ v ] = 정점 v와 v에서 가장 먼 서브트리 외 정점 사이의 거리

 따라서 위의 두 값을 구하면, far [ v ] = max( dist [ v ], parent_far [ v ] ) 로 구할 수 있습니다.

1) dist [ v ] = 정점 v와 v에서 가장 먼 서브트리 내 정점 사이의 거리

 dist[v]는 간단한 트리 동적 계획법으로 구할 수 있습니다. dist [ v ] = MAX( dist[자식 노드] + 1 ) 이기 때문입니다.

#include<stdio.h>
#include<algorithm>
#include<vector>
 
using namespace::std;
 
 
const int N=100005;
 
int n,m,q;
int dist[N];
int check[N];
int chk_n=1;
vector<int> ad[N];
 
void dist_dfs(int now){
    check[now] = chk_n;
    dist[now]=0;
    for(auto next: ad[now]){
        if(check[next]!=chk_n){
            dist_dfs(next);
            dist[now] = max(dist[now], dist[next]+1);
        }
    }
}
 
int main() {
    /*
    간선 (i,j) 를 ad[i].push_back(j) 로 입력.
    */
    dist_dfs(1);
    return 0;
}

2) parent_far [ v ] = 정점 v와 v에서 가장 먼 서브트리 외 정점 사이의 거리

 parent_far[v]를 구하기 위해, 서브트리 외 정점들을 2가지 그룹으로 나누어야 합니다. 나누는 기준은 v 의 부모 정점의 서브트리에 있는 지의 여부입니다. 아래와 같은 트리의 정점 4 에서 parent_far [ 4 ] 를 구할 때의 예시는 아래와 같습니다. 4의 부모는 2이기 때문에, 2의 서브트리에 있는지의 여부를 따지게 되지요.

경우 1 : 부모 정점의 서브트리에 없는 경우는 parent_far [ 부모 정점 ] + 1 로 구할 수 있습니다.

             (v -> 부모 -> parent_far[부모] 로 가는 경로가 최장임을 증명할 수 있습니다.)

경우 2 : 부모 정점의 서브트리에 있는 경우는 대개 dist [ 부모 정점 ] + 1 로 구할 수 있습니다.

             (v -> 부모 -> dist[부모] 로 가는 경로가 최장인 경우가 많습니다.)

 다만 dist [ 부모 노드 ] 에서 선택된 부모 노드에서 가장 먼 정점이 지금 정점 v의 서브 트리에 있을 경우엔 dist [ 부모 정점 ] + 1 가 답이 아닙니다. 아래에 그 예시가 있습니다.

 이 경우를 위해서 dist [ 부모 노드 ] 연산 시 dist [ v ] + 1 을 제외하고 연산한 결과가 필요하게 됩니다. 이를 위해 dist2 [ 정점 ] = dist [ 자식 노드 ] +1 중에서 2 번째로 큰 값을 미리 저장하면, 경우 2의 답을 dist2 [ 부모 정점 ] + 1 로 구할 수 있습니다. (v-> 부모-> dist2[부모] 로 가는 경로가 최장입니다.)

 아래의 코드를 통해 parent_far 와 dist2 를 구하는 방법을 자세히 알 수 있습니다.

3. 방법 1 코드

//code by RiKang, weeklyps.com
#include<stdio.h>
#include<algorithm>
#include<vector>
using namespace::std;
const int N=100005;
int n,m,q;
int dist[N];
int dist2[N];
int far[N];
int check[N];
int chk_n=1;
vector<int> ad[N];
 
void dist_dfs(int now){
    check[now] = chk_n;
    dist[now] = dist2[now] = 0;
    for(auto next: ad[now])
        if(check[next]!=chk_n){
            dist_dfs(next);
            if(dist[next]+1 > dist[now]){
                dist2[now] = dist[now];  // dist2 에 두 번째로 큰 값을 저장
                dist[now] = dist[next]+1;
            }
            else if(dist[next]+1 > dist2[now]){
                dist2[now] = dist[next]+1;// dist2 에 두 번째로 큰 값을 저장
            }
        }
}
void far_dfs(int now, int parent_far){
    check[now] = chk_n;
    far[now] = max(dist[now], parent_far);
    for(auto next: ad[now])
        if(check[next]!=chk_n){
            if(dist[now] == dist[next]+1) // next->now->dist[now]가 안되는 상황
                far_dfs(next, max(parent_far+1, dist2[now]+1));
            else
                far_dfs(next, max(parent_far+1, dist[now]+1));
        }
}
int main() {
    /*
    간선 (i,j) 를 ad[i].push_back(j) 로 입력.
    */
    dist_dfs(1);
    // 여기에서도, 트리의 지름 = MAX(dist[v]+dist2[v]); 로 계산 가능.
    chk_n++;
    far_dfs(1, 0);
    // 트리의 지름 = MAX(far[v]);
    return 0;
}
 

 사실 far 배열은 필요 없고 트리의 지름만 구하고 싶다면 dist[v]+dist2[v] 의 최댓값을 통해 찾을 수도 있습니다. v가 경로에 속하면서 v 의 부모는 경로에 속하지 않는 경로 중, 가장 긴 경로가 dist[v]+dist2[v] 이기 때문입니다. 어떤 경로든지, v가 경로에 속하면서 v 의 부모는 경로에 속하지 않는 v 는 하나 가지고 있으므로, dist[v]+dist2[v] 의 최댓값도 모든 경우의 수를 고려한 값이 됩니다.

4. 방법 2 - 탐욕법

 방법 1에선 다소 복잡한 방법을 사용했지만, 사실 트리의 지름은 아래의 간단한 그리디 패러다임의 알고리즘을 따라가면 구할 수 있습니다.

step 1) 트리에서 임의의 정점 v1 에서 가장 먼 정점 v2를 찾는다.

step 2) v2와 v2 에서 가장 먼 정점 v3 사이의 거리를 찾는다.

step 3) 이 v2 와 v3 사이의 거리를 구하면 그게 트리의 지름이 된다.

 트리에서 v1 이 주어졌을 때, 가장 먼 정점 v2 를 찾는 건 트리 DP로 구할 수 있습니다. 설명은 코드 아래에 있습니다.

//code by RiKang, weeklyps.com
#include<stdio.h>
#include<algorithm>
#include<vector>
 
using namespace::std;
 
const int N=100005;
int n;
int dist[N];
int check[N];
int chk_n=1;
vector<int> ad[N];
 
int dist_dfs(int now){
    check[now] = chk_n;
    int ret=now;
    dist[now]=0;
    for(auto next: ad[now]){
        if(check[next]!=chk_n){
            int ret_next = dist_dfs(next);
            if(dist[next]+1 > dist[now]){
                dist[now] = dist[next]+1;
                ret = ret_next;
            }
        }
    }
    return ret;
}

 dist [ v ] = 정점 v와 v에서 가장 먼 서브트리 내 정점 사이의 거리

 위의 배열은 간단한 트리 동적 계획법으로 구할 수 있습니다. dist [ v ] = MAX( dist[자식 노드] + 1 ) 이기 때문입니다. 따라서 v1에서 가장 먼 거리를 찾고 싶다면, v1을 루트 노드로 생각하고 dist_dfs(v1); 을 호출하면 됩니다. 그러면 dist[v1] 이 v1에서 가장 먼 거리가 되겠지요.

 그리고 위와 같이 dist_dfs() 가 목적지 노드를 리턴하도록 만들면, 결과적으로 루트(v1) 에서 가장 먼 정점이 어디인지를 리턴받을 수 있습니다.

 따라서 아래의 코드를 통해 step 1,2,3 을 추가로 진행하면 트리의 지름을 구할 수 있습니다.

5. 방법 2 코드

//code by RiKang, weeklyps.com
#include<stdio.h>
#include<algorithm>
#include<vector>
 
using namespace::std;
 
const int N=100005;
int n;
int dist[N];
int check[N];
int chk_n=1;
vector<int> ad[N];
 
int dist_dfs(int now){
    check[now] = chk_n;
    int ret=now;
    dist[now]=0;
    for(auto next: ad[now]){
        if(check[next]!=chk_n){
            int ret_next = dist_dfs(next);
            if(dist[next]+1 > dist[now]){
                dist[now] = dist[next]+1;
                ret = ret_next;
            }
        }
    }
    return ret;
}
int main() {
    /* 간선 입력 */
    int v2 = dist_dfs(1);
    chk_n++;
    int v3 = dist_dfs(v2);
    printf("%d",dist[v3]); // 지름
    return 0;
}

6. 방법 2 의 정당성

step 1) 트리에서 임의의 정점 v1 에서 가장 먼 정점 v2를 찾는다.

step 2) v2와 v2 에서 가장 먼 정점 v3 사이의 거리를 찾는다.

step 3) 이 v2 와 v3 사이의 거리를 구하면 그게 트리의 지름이 된다.

 위의 과정이 트리의 지름을 찾음을 증명해보겠습니다. 우선 어떤 정점 쌍 (w1, w2) 가 트리의 지름이라고 하면, 경우의 수는 3가지가 있습니다.

경우 1) v1 == w1 || v1==w2

경우 2) v2 == w1 || v2==w2

경우 3) v1, v2, w1, w2 가 모두 다른 경우

 d(v1, v2) = v1과 v2사이의 거리, 라고 지정하고 이 3가지 경우에 대해 따로따로 증명해보겠습니다.

경우 1) v1 == w1 || v1==w2

 w1, w2 중 하나가 v1 이므로 'd(w1,w2) <= v1에서 가장 먼 거리'가 성립합니다.

 그리고 v1,v2,w1,w2의 정의에 의해 'd(w1,w2)=지름'이고, 'd(v1,v2)=v1에서 가장 먼 거리'입니다.

 -> 지름 = d(w1,w2) <= d(v1,v2) <= d(v2,v3) 으로 증명 끝.

경우 2) v2 == w1 || v2==w2

 지름 = d(w1,w2) <= v2에서 가장 먼 거리 = d(v2,v3) 으로 증명 끝.

경우 3) v1, v2, w1, w2 가 모두 다른 경우

3-a). 4 정점이 아래와 같이 연결되어 있을 경우

v1에서 가장 먼 거리 = d(v1,v2) >= d(v1,w2)

-> d(v1,A) + d(A,v2) >= d(v1,A) + d(A,w2)

-> d(A,v2) >= d(A,w2)

-> d(w1,A) + d(A,v2) >= d(w1,A) + d(A,w2)

-> d(w1,v2) >= d(w1,w2)

 따라서 (w1,v2) 도 트리의 지름이 됩니다. -> (w1,v2)가 트리의 지름일 경우, (v1, v2)도 지름인 것은 경우 2에서 증명하였습니다.

3-b). 4 정점이 아래와 같이 연결되어 있을 경우

3-a) 와 대동소이합니다.

v1에서 가장 먼 거리 = d(v1,v2) >= d(v1,w2)

-> d(v1,B) + d(B,v2) >= d(v1,B) + d(B,w2)

-> d(B,v2) >= d(B,w2)

-> d(w1,B) + d(B,v2) >= d(w1,B) + d(B,w2)

-> d(w1,v2) >= d(w1,w2)

 따라서 (w1,v2) 도 트리의 지름이 됩니다. -> (w1,v2)가 트리의 지름일 경우, (v1, v2)도 지름인 것은 경우 2에서 증명하였습니다.

7. 방법 2 로 far 배열 찾기

 방법 2에서도 far [ v ] = 정점 v와 v에서 가장 먼 정점 사이의 거리을 모든 정점에 대해 쉽게 구할 수 있습니다. d(정점, 정점) 이 두 정점 사이의 거리라고 하면 d(v2, v3)가 트리의 지름이며, far [ v ] = MAX( d(v,v2), d(v,v3) ) 가 성립하기 때문입니다.

증명 ) 위의 방법 2 정당성 증명과 대동소이 하므로, 간략하게 설명하겠습니다.

 우선 v에서 가장 먼 정점을 w 라고 하겠습니다.

경우 1) v, w, v2, v3가 아래와 같이 연결된 경우

(v2,v3) 가 지름이기 때문에 d(v,v3) >= d(v,w) 이므로 문제될 것 없습니다.

경우 2) v, w, v2, v3가 아래와 같이 연결된 경우

 (v2,v3) 가 지름이기 때문에 d(v,v3) >= d(v,w) 이므로 문제될 것 없습니다.

 far [ v ] = MAX( d(v,v2), d(v,v3) ) 를 구하는 건, dist2 [ v ] = d(v2,v), dist3 [ v ] = d(v3,v) 의 배열을 트리 순회 과정에서 간단하게 구할 수 있기 때문에 여기까지 읽으신 분들은 당연히 할 수 있을 거라고 믿고(?) 생략하겠습니다.

8. 문제

(0) [BOJ 1967] 트리의 지름

//code by RiKang, weeklyps.com
#include<stdio.h>
#include<algorithm>
#include<vector>
 
using namespace::std;
 
const int N=100005;
int n;
int dist[N], dist2[N], far[N];
int check[N];
int chk_n=1;
vector<int> ad[N];
vector<int> cost[N];
void dist_dfs(int now){
    check[now] = chk_n;
    dist[now] = dist2[now] = 0;
    auto next_cost = cost[now].begin();
    for(auto next: ad[now]){
        if(check[next]!=chk_n){
            dist_dfs(next);
            if(dist[next]+*next_cost > dist[now]){
                dist2[now] = dist[now];
                dist[now] = dist[next]+*next_cost;
            }
            else if(dist[next]+*next_cost > dist2[now]){
                dist2[now] = dist[next]+*next_cost;
            }
        }
        next_cost++;
    }
}
void far_dfs(int now, int parent_far){
    check[now] = chk_n;
    far[now] = max(dist[now], parent_far);
    auto next_cost = cost[now].begin();
    for(auto next: ad[now]){
        if(check[next]!=chk_n){
            if(dist[now] == dist[next]+*next_cost)
                far_dfs(next, max(parent_far+*next_cost, dist2[now]+*next_cost));
            else
                far_dfs(next, max(parent_far+*next_cost, dist[now]+*next_cost));
        }
        next_cost++;
    }
}
int main() {
    int n;
    scanf("%d",&n);
    for(int i=1; i<n; i++){
        int v1,v2,c;
        scanf("%d %d %d",&v1,&v2,&c);
        ad[v1].push_back(v2);
        cost[v1].push_back(c);
        ad[v2].push_back(v1);
        cost[v2].push_back(c);
    }
    dist_dfs(1);
    int ans1 = 0;
    for(int i=1; i<=n; i++)
        ans1 = max(ans1, dist[i]+dist2[i]); // 이것도 정답
    chk_n++;
    far_dfs(1, 0);
    int ans = 0;
    for(int i=1; i<=n; i++)
        ans = max(ans, far[i]); // 이것도 정답
    if(ans==ans1)
        printf("%d",ans);
    return 0;
}

//code by RiKang, weeklyps.com
#include<stdio.h>
#include<algorithm>
#include<vector>
 
using namespace::std;
 
const int N=100005;
int n;
int dist[N];
int check[N];
int chk_n=1;
vector<int> ad[N];
vector<int> cost[N];
 
int dist_dfs(int now){
    check[now] = chk_n;
    int ret=now;
    dist[now]=0;
    auto next_cost = cost[now].begin();
    for(auto next: ad[now]){
        if(check[next]!=chk_n){
            int ret_next = dist_dfs(next);
            if(dist[next]+*next_cost > dist[now]){
                dist[now] = dist[next]+*next_cost;
                ret = ret_next;
            }
        }
        next_cost++;
    }
    return ret;
}
int main() {
    int n;
    scanf("%d",&n);
    for(int i=1; i<n; i++){
        int u,v,c;
        scanf("%d %d %d",&u,&v,&c);
        ad[u].push_back(v);
        cost[u].push_back(c);
        ad[v].push_back(u);
        cost[v].push_back(c);
    }
    int v2 = dist_dfs(1);
    chk_n++;
    int v3 = dist_dfs(v2);
    printf("%d",dist[v2]); // 지름
    return 0;
}

(1) [CF 411 D] Expected diameter of a tree

//code by RiKang, weeklyps.com
#include<stdio.h>
#include<algorithm>
#include<vector>
#include<map>
#define PII pair<int, int>
 
using namespace::std;
 
const int N=100005;
 
struct union_find_tree{
    vector<int> parent;
    vector<int> height;
    vector<double> group_size;
 
    void init(int uf_n){
        parent.clear(), height.clear();
        group_size.clear();
        for(int i=0; i<uf_n; i++){
            parent.push_back(i);
            height.push_back(0);
            group_size.push_back(1);
        }
    }
 
    int get_root(int now){
        if(parent[now]==now) return now;
        return parent[now] = get_root(parent[now]);
    }
 
    void merge(int v1,int v2){
        v1=get_root(v1), v2=get_root(v2);
        if(v1==v2) return;
        if(height[v1] < height[v2]) swap(v1,v2);
        parent[v2]=v1;
        group_size[v1]+=group_size[v2];
        if(height[v1]==height[v2])
            height[v1]++;
    }
}uf;
 
int n,m,q;
int dist[N];
int dist2[N];
int far[N];
map<PII, double> ma;
int check[N];
int chk_n=1;
vector<int> ad[N];
vector<int> cnt[N];
 
void dist_dfs(int now){
    check[now] = chk_n;
    dist[now] = dist2[now] = 0;
    for(auto next: ad[now])
        if(check[next]!=chk_n){
            dist_dfs(next);
            if(dist[next]+1 > dist[now]){
                dist2[now] = dist[now];
                dist[now] = dist[next]+1;
            }
            else if(dist[next]+1 > dist2[now]){
                dist2[now] = dist[next]+1;
            }
        }
}
 
void far_dfs(int now, int parent_far){
    check[now] = chk_n;
    far[now] = max(dist[now], parent_far);
    for(auto next: ad[now])
        if(check[next]!=chk_n){
            if(dist[now] == dist[next]+1){
                far_dfs(next, max(parent_far+1, dist2[now]+1));
            }
            else
                far_dfs(next, max(parent_far+1, dist[now]+1));
        }
}
 
double get_ans(int v1, int v2){
    double ret = (double)cnt[v1].size()-1;
    double op = 0.0, sum = 0.0, di = uf.group_size[v1] * uf.group_size[v2];
    for(int i = 0; i<cnt[v2].size(); i++){
        sum+=cnt[v1][cnt[v1].size()-1-i];
        op+=sum/di;
        ret += op * cnt[v2][i];
    }
    return ret;
}
 
int main(){
    scanf("%d %d %d",&n,&m,&q);
    uf.init(N);
    while(m--){
        int v1, v2;
        scanf("%d %d",&v1,&v2);
        ad[v1].push_back(v2);
        ad[v2].push_back(v1);
        uf.merge(v1,v2);
    }
    for(int i=1; i<=n; i++)
        if(uf.get_root(i) == i){
            chk_n++;
            dist_dfs(i);
            chk_n++;
            far_dfs(i,0);
        }
    for(int i=1; i<=n; i++){
        int s = uf.get_root(i);
        if((int)cnt[s].size()-1<far[i])
            cnt[s].resize(far[i]+1,0);
        cnt[s][far[i]]++;
    }
    while(q--){
        int o1, o2;
        scanf("%d %d",&o1,&o2);
        o1 = uf.get_root(o1);
        o2 = uf.get_root(o2);
        if(o1==o2){
            printf("-1\n");
            continue;
        }
        if(cnt[o1].size()<cnt[o2].size() || (cnt[o1].size()==cnt[o2].size() && o1<o2))
            swap(o1,o2);
        if(ma.find(PII(o1,o2))==ma.end())
            ma[PII(o1,o2)] = get_ans(o1,o2);
        printf("%.9lf\n",ma[PII(o1,o2)]);
    }
    return 0;
}

※ 본문의 코드들은 C++ 17 환경에서 작성되었습니다.

※ 문의 : rikang93@gmail.com