백준 1167번 트리의 지름 (Java)

문제

백준 1167번 트리의 지름 / 골드2
트리의 지름이란, 트리에서 임의의 두 점 사이의 거리 중 가장 긴 것을 말한다. 트리의 지름을 구하는 프로그램을 작성하시오.

입력

트리가 입력으로 주어진다. 먼저 첫 번째 줄에서는 트리의 정점의 개수 V가 주어지고 (2 ≤ V ≤ 100,000)둘째 줄부터 V개의 줄에 걸쳐 간선의 정보가 다음과 같이 주어진다. 정점 번호는 1부터 V까지 매겨져 있다.

먼저 정점 번호가 주어지고, 이어서 연결된 간선의 정보를 의미하는 정수가 두 개씩 주어지는데, 하나는 정점번호, 다른 하나는 그 정점까지의 거리이다. 예를 들어 네 번째 줄의 경우 정점 3은 정점 1과 거리가 2인 간선으로 연결되어 있고, 정점 4와는 거리가 3인 간선으로 연결되어 있는 것을 보여준다. 각 줄의 마지막에는 -1이 입력으로 주어진다. 주어지는 거리는 모두 10,000 이하의 자연수이다.

풀이

트리의 특성을 활용해 두 정점 사이의 최대 거리를 구하는 문제이다. 처음에는 BFS와 백트래킹을 사용해 모든 정점에서 그 외의 정점으로 가는 비용의 최댓값을 계산했다. 하지만 정점의 개수 V의 최대 개수는 100,000이기 때문에 시간 복잡도 O(V²)로 당연하게도 시간 초과가 발생했다. 그럼 어떻게 문제를 해결해야 하나 고민 끝에 트리의 특성을 탐구했다.

트리는 구조상 사이클이 만들어지지 않는다. 그렇기 때문에 한 정점에서 다른 정점으로 가는 경로는 유일하다. 또한, 한 정점에서 가장 먼 정점으로 가는 경로는 문제에서 구하고자 하는 최대 거리의 경로와 겹치게 된다. 그 이유는 임의의 정점을 루트라고 생각해 보자. 루트에서 왼쪽의 가장 먼 정점인 리프 노드a(임시)로 가는 경로와 오른쪽의 가장 먼 정점인 리프 노드b(임시)로 가는 경로를 합하면 최대 거리가 되는데, 트리의 모양을 생각하면 어떤 노드를 선택해 최대 거리를 구하더라도 일전의 a 혹은 b 중 하나를 포함할 수밖에 없다.

이제 코드를 보자. 우선, 임의의 정점start를 잡아 BFS를 돌려, start로 시작해 모든 노드로 가는 최대 거리를 저장한 distance배열을 구했다. 그리고 최대 거리를 저장한 정점의 번호maxIdx를 구하게 되면, maxIdx는 리프 노드이며 구하려고 하는 최대 거리의 경로에 무조건 포함이 된다. 그렇게 되면 maxIdx를 필두로 한 BFS를 한 번 더 돌리게 되면 최대 경로의 값이 자연스럽게 구해지게 된다.

코드

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int V = Integer.parseInt(br.readLine());

        List<Node>[] tree = new ArrayList[V + 1];
        for (int i = 1; i <= V; i++) {
            tree[i] = new ArrayList<Node>();
        }

        StringTokenizer st;
        for (int i = 0; i < V; i++) {
            st = new StringTokenizer(br.readLine());
            int from = Integer.parseInt(st.nextToken());
            while (st.hasMoreTokens()) {
                int to = Integer.parseInt(st.nextToken());
                if (to == -1) {
                    break;
                }
                int dis = Integer.parseInt(st.nextToken());
                tree[from].add(new Node(to, dis));
            }
        }

        int[] distance = new int[V + 1];
        bfs(tree, distance, 1);
        int maxIdx = 1;
        for (int i = 2; i <= V; i++) {
            if (distance[maxIdx] < distance[i]) maxIdx = i;
        }
        distance = new int[V + 1];
        bfs(tree, distance, maxIdx);

        Arrays.sort(distance);
        System.out.println(distance[V]);
    }

    static void bfs(List<Node>[] tree, int[] distance, int start) {
        Queue<Integer> queue = new LinkedList<Integer>();
        boolean[] visit = new boolean[tree.length];
        queue.offer(start);
        visit[start] = true;

        while (!queue.isEmpty()) {
            int now = queue.poll();
            for (Node next : tree[now]) {
                int nextIdx = next.to;
                int nextDis = distance[now] + next.dis;
                if (visit[nextIdx]) continue;
                visit[nextIdx] = true;
                queue.offer(nextIdx);
                distance[nextIdx] = nextDis;
            }
        }
    }
}

class Node {
    int to;
    int dis;

    Node(int to, int dis) {
        this.to = to;
        this.dis = dis;
    }
}