[백준] 1167번 - 트리의 지름 (Gold 2)

문제 링크

[백준] 1167번 - 트리의 지름 (Gold 2)

문제 설명

트리와 각 간선의 가중치가 입력으로 주어진다.
트리의 두 점 사이의 거리 중 가장 긴 거리를 출력한다.

정답 코드 및 설명

사이클이 없는 연결 그래프는 임의의 노드를 루트 노드로 취급한 트리로 생각할 수 있다.
1번 노드를 루트 노드로 취급하자.
트리의 지름은 다음의 알고리즘을 통해 구할 수 있다.

1. 루트 노드에서 가장 먼 노드를 찾는다.
2. 1에서 찾은 노드로부터 가장 먼 노드를 찾는다.
3. 1에서 찾은 노드와 2에서 찾은 노드 사이의 거리가 트리의 지름이다.

알고리즘이 왜 성립하는지는 아래에서 다루기로 하고, 정답 코드를 먼저 보자.

정답 코드

maxDist(int index) 메소드가 index 노드로 부터 가장 먼 노드를 찾아주는 부분이다.
index 노드로부터의 거리를 계산할때는 index 노드를 루트 노드로 생각하고 DFS를 적용했다.
루트 노드 - 자식 노드 거리 = 루트 노드 - 부모 노드 거리 + 부모 노드 - 자식 노드 거리임을 활용했다.
BFS로도 수행이 가능하기는 하나, 부모 노드도 같이 기록해야하므로 DFS를 쓰는 것이 보다 편리하다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.StringTokenizer;

public class BOJ1167 {
    int n;
    ArrayList<Node>[] adj;

    class Node {
        int node, dist;

        Node(int node, int dist) {
            this.node = node;
            this.dist = dist;
        }
    }

    void makeGraph() throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        n = Integer.parseInt(br.readLine());
        adj = new ArrayList[n];
        for (int i = 0; i < n; i++)
            adj[i] = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int start = Integer.parseInt(st.nextToken()) - 1;
            while (true) {
                int end = Integer.parseInt(st.nextToken()) - 1;
                if (end == -2) break;
                int weight = Integer.parseInt(st.nextToken());
                adj[start].add(new Node(end, weight));
            }
        }
    }

    Node maxDist(int index) {
        int[] dist = new int[n];
        Arrays.fill(dist, -1);
        dist[index] = 0;
        dfs(index, dist);
        int maxIndex = 0;
        int maxDist = 0;
        for (int i = 0; i < n; i++) {
            if (dist[i] > maxDist) {
                maxDist = dist[i];
                maxIndex = i;
            }
        }
        return new Node(maxIndex, maxDist);
    }

    void dfs(int index, int[] dist) {
        int size = adj[index].size();
        for (int i = 0; i < size; i++) {
            Node next = adj[index].get(i);
            if (dist[next.node] == -1) {
                dist[next.node] = dist[index] + next.dist;
                dfs(next.node, dist);
            }
        }
    }

    void solution() throws IOException {
        makeGraph();
        int a = maxDist(0).node;
        int radius = maxDist(a).dist;
        System.out.println(radius);
    }

    public static void main(String[] args) throws IOException {
        new BOJ1167().solution();
    }
}

알고리즘이 성립하는 이유

위 알고리즘을 통해 트리의 지름을 구할 수 있음을 증명하려면 다음의 두 정리가 필요하다.

[정리 1] 트리에서 가장 먼 두 노드 중 적어도 하나는 루트 노드에서 가장 먼 노드이다.
[정리 2] 루트 노드에서 가장 먼 노드 a와 a에서 가장 먼 노드 b 사이의 거리는 트리의 지름과 같다.

증명에 앞서, 용어 하나를 정의하자.
최소 공통 조상(Lowest Common Ancestor, LCA)이란, 두 노드의 공통된 조상 중에서 가장 가까운 조상이다.

• $LCA(x, y)$는 x번 노드와 y번 노드의 최소 공통 조상을 뜻한다.
• $d(x,y)$는 x번 노드와 y번 노드의 거리이다.

정리 1의 증명

트리에서 가장 먼 두 노드가 a번 노드와 b번 노드라고 하자.
루트 노드에서 가장 먼 노드를 c번 노드라고 두자.

$d(root, c) > d(root, a)$, $d(root, c) > d(root, b)$라고 가정하자.
이 때 가능한 경우는 총 3가지가 있다.

1. $LCA(a, b) = x$, $LCA(a, c) = LCA(b, c) = y$
   • y는 x의 조상이거나 y = x
   • $d(a, c) = d(y, c) + d(a, y)$
   • $d(y, c) = d(root, c) - d(root, y) > d(root, b) - d(root, x) = d(x, b)$
   • $d(a, y) = d(x, y) + d(a, x) \geq d(a, x)$
   • 종합하면, $d(a, c) > d(x, b) + d(a, x) = d(a, b)$
   • a번 노드와 c번 노드 사이의 거리가 더 멀어지므로 모순
2. $LCA(a, c) = x$, $LCA(a, b) = LCA(b, c) = y \neq x$
   • y는 x의 조상
   • $d(b, c) = d(y, c) + d(b, y)$
   • $d(y, c) = d(root, c) - d(root, y) > d(root, a) - d(root, y) = d(y, a)$
   • 종합하면, $d(b, c) > d(y, a) + d(b, y) = d(b, a)$
   • b번 노드와 c번 노드 사이의 거리가 더 멀어지므로 모순
3. $LCA(b, c) = x$, $LCA(a, b) = LCA(a, c) = y \neq x$
   • 사실상 3번 경우와 동일하다. a와 b만 맞바꿔주면 된다.

3가지 경우 모두 모순이 발견되었으므로, $d(root, c) = d(root, a)$ 또는 $d(root, c) = d(root, b)$이다.

정리 2의 증명

루트 노드에서 가장 먼 노드를 a번 노드라고 하고, a번 노드에서 가장 먼 노드를 b번 노드라고 하자.
두 노드 c, d가 트리에서 가장 먼 두 노드이고, $d(c, d) > d(a, b)$라고 가정하자.

1. $LCA(c, d) = x$, $LCA(a, c) = LCA(a, d) = y \neq x$
   • y는 x의 조상
   • [정리 1]에 의해, $d(root, a) = d(root, c)$ 또는 $d(root, a) = d(root, d)$
   • $d(root, a) = d(root, c)$인 경우
   • $d(a, d) = d(y, d) + d(a, y) = d(y, x) + d(x, d) + d(a, y) > d(x, d) + d(a, y)$
   • $d(c, d) = d(x, d) + d(c, x)$
   • $d(a, y) = d(root, a) - d(root, y) > d(root, c) - d(root, x) = d(c, x)$
   • $d(a, b) \geq d(a, d) > d(x, d) + d(a, y) > d(x, d) + d(c, x) = d(c, d)$
   • 가정 $d(c, d) > d(a, b)$에 모순
   • $d(root, a) = d(root, d)$인 경우는 위에서 c와 d만 맞바꿔주면 된다.
2. $LCA(a, d) = x$, $LCA(a, c) = LCA(c, d) = y$
   • y는 x의 조상
   • $d(c, a) = d(y, a) + d(c, y)$
   • $d(y, a) = d(root, a) - d(root, y) > d(root, d) - d(root, y) = d(y, d)$
   • $d(a, b) \geq d(c, a) > d(y, d) + d(c, y) = d(c, d)$
   • 가정 $d(c, d) > d(a, b)$에 모순
3. $LCA(a, c) = x$, $LCA(a, d) = LCA(c, d) = y$
   • 사실상 2번 경우와 동일하다. c와 d만 맞바꿔주면 된다.

3가지 경우 모두 모순이 발견되었으므로, $d(c, d) \leq d(a, b)$이다.
c와 d의 정의에 의해 $d(c, d) \geq d(a, b)$이므로, $d(a, b) = d(c, d)$이다.