(Java) 백준 1167

백준 1167

시간 제한	메모리 제한	제출	정답	맞힌 사람	정답 비율
2 초	256 MB	69750	24076	17533	34.169%

문제

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

입력

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

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

출력

첫째 줄에 트리의 지름을 출력한다.

---

답안

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class BOJ1167 {
    static class GraphNode {
        int endVertex;
        int weight;

        public GraphNode(int endVertex, int weight) {
            this.endVertex = endVertex;
            this.weight = weight;
        }
    }

    static int answer = 0;

    public static void main(String[] args) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
        int vertexCount = Integer.parseInt(reader.readLine());

        Hashtable<Integer, ArrayList<GraphNode>> graph = new Hashtable<>();
        for (int i = 0; i < vertexCount; i++) {
            graph.put(i, new ArrayList<>());
        }
        for (int i = 0; i < vertexCount; i++) {
            StringTokenizer tokenizer = new StringTokenizer(reader.readLine());
            int startVertex = Integer.parseInt(tokenizer.nextToken()) - 1;
            int neighborsCount = (tokenizer.countTokens() - 1) / 2;
            for (int j = 0; j < neighborsCount; j++) {
                int endVertex = Integer.parseInt(tokenizer.nextToken()) - 1;
                int weight = Integer.parseInt(tokenizer.nextToken());
                graph.get(startVertex).add(new GraphNode(endVertex, weight));
                graph.get(endVertex).add(new GraphNode(startVertex, weight));
            }
        }

        boolean[] visited = new boolean[vertexCount];

        dfs(0, graph, visited);

        System.out.println(answer);
    }

    static int dfs(int start, Hashtable<Integer, ArrayList<GraphNode>> graph, boolean[] visited) {
        visited[start] = true;
        int[] topTwoLen = new int[2];

        for (GraphNode neighbor : graph.get(start)) {
            if (visited[neighbor.endVertex]) {
                continue;
            }
            int childLen = dfs(neighbor.endVertex, graph, visited) + neighbor.weight;
            if (childLen >= topTwoLen[0]) {
                topTwoLen[1] = topTwoLen[0];
                topTwoLen[0] = childLen;
            } else if (childLen >= topTwoLen[1]) {
                topTwoLen[1] = childLen;
            }
            answer = Math.max(answer, topTwoLen[0] + topTwoLen[1]);
        }

        return topTwoLen[0];
    }
}

풀이

DFS를 이용해 풀었다.

인터넷 블로그 등에서 많이 활용되는 방법은 DFS를 두 번 하는 것이지만, 나는 DFS를 한 번만 하는 방법으로 문제를 해결했다.

1번만 DFS를 했을 때 문제점은, 트리의 지름 경로가, DFS를 시작하는 임의의 노드를 지나지 않는 경우, 지름을 구할 수 없는 것이다.

그래서, 이를 보완하기 위해 각 노드를 루트로 하는 서브트리에서 지름이 존재하는지 확인하도록 했다.

각 노드를 루트로 하는 서브트리에서, 자신의 자식 노드로의 경로 중 길이가 가장 긴 것 2개를 더한 값을 answer과 비교해 업데이트했고, 이를 모든 서브트리에 대해 적용했다.

가장 중요한 코드 부분은 다음과 같다.

// 자식의 길이 = 그 자식의 노드부터 leaf 노드까지의 거리 + 현재 노드와 자식 노드 사이의 거리
int childLen = dfs(neighbor.endVertex, graph, visited) + neighbor.weight;

// 자식의 길이가 가장 긴 것 2개를 구함
if (childLen >= topTwoLen[0]) {  // > 가 아닌 >=를 사용해야 함에 유의
    topTwoLen[1] = topTwoLen[0];
    topTwoLen[0] = childLen;
} else if (childLen >= topTwoLen[1]) {
    topTwoLen[1] = childLen;
}

// 그 둘을 더한 것이 전체 트리의 지름인지 확인
answer = Math.max(answer, topTwoLen[0] + topTwoLen[1]);

참고로, 이 풀이는 어느 노드에서 DFS를 시작하든 상관 없다. 트리의 지름 경로가 루트를 지나든 지나지 않든 알고리즘이 동작하므로, 임의의 노드를 루트로 가정해도 되기 때문이다.