[Gold II] 트리의 지름 - 1167

문제 링크

분류

깊이 우선 탐색, 그래프 이론, 그래프 탐색, 트리

문제 설명

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

입력

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

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

출력

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

풀이

트리 구조로 정점들이 이루어져있을때, 정점간의 최대거리를 구하는 문제였다. 임의의 한 점을 트리의 root로 설정하고, 정점의 서브트리를 탐색하는 재귀 함수를 설정하였다. 재귀 함수는 정점의 자녀 중 가장 길이가 먼 길이를 리턴하게했으며, root를 방문하지 않고 최대 거리가 나오는 경우가 있을 수 있기에, 재귀 함수 내부에서 서브 트리 거리의 합도 계산함으로서 답을 구할 수 있었다.

import java.io.*;
import java.util.*;

class Main {

    public static final BufferedReader BR = new BufferedReader(new InputStreamReader(System.in));
    public static final BufferedWriter BW = new BufferedWriter(new OutputStreamWriter(System.out));
    int V;
    long answer;
    HashMap<Integer, Node> nodeHashMap;
    HashSet<Integer> visited;

    public static void main(String[] args) {
        Main main = new Main();
        try {
            main.init();
            main.solution();
        } catch (Exception e) {
            System.out.println("exception during I/O");
        }

    }

    void init() throws Exception {
        V = Integer.parseInt(BR.readLine());
        answer = 0L;
        nodeHashMap = new HashMap<>();
        visited = new HashSet<>();

        for (int i = 0; i < V; i++) {
            int[] inputArray = Arrays.stream(BR.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
            int from = inputArray[0];

            if (!nodeHashMap.containsKey(from)) {
                Node newNode = new Node(from);
                nodeHashMap.put(from, newNode);
            }

            Node node = nodeHashMap.get(from);

            for (int j = 1; j + 1 < inputArray.length; j += 2) {
                node.neighbor.put(inputArray[j], inputArray[j + 1]);
            }

        }

    }

    void solution() throws Exception {
        Node root = nodeHashMap.get(1);
        visited.add(1);
        long fromRoot = travel(root);
        if (fromRoot > answer) {
            answer = fromRoot;
        }
        System.out.println(answer);
    }

    long travel(Node parent) {
        HashMap<Integer, Integer> neighbors = parent.getNeighbor();
        PriorityQueue<Long> maxSubTreeDistance = new PriorityQueue<>(Collections.reverseOrder());

        for (Integer child : neighbors.keySet()) {
            if (!visited.contains(child)) {
                visited.add(child);
                Node nextNode = nodeHashMap.get(child);
                long travelDistance = travel(nextNode);
                travelDistance += neighbors.get(child);
                maxSubTreeDistance.add(travelDistance);
            }
        }

        if (maxSubTreeDistance.isEmpty()) {
            return 0L;
        } else {
            long toReturn = maxSubTreeDistance.peek();
            if (maxSubTreeDistance.size() >= 2) {
                long subTreeDistance1 = maxSubTreeDistance.remove();
                long subTreeDistance2 = maxSubTreeDistance.remove();
                if (subTreeDistance1 + subTreeDistance2 > answer) {
                    answer = subTreeDistance1 + subTreeDistance2;
                }
            }
            return toReturn;
        }
    }

    static class Node {
        int number;
        HashMap<Integer, Integer> neighbor;

        public Node(int number) {
            this.number = number;
            this.neighbor = new HashMap<>();
        }

        public int getNumber() {
            return number;
        }

        public HashMap<Integer, Integer> getNeighbor() {
            return neighbor;
        }
    }


}