[트리의 지름 구하기] 백준 1167 (JAVA) - dfs / 다익스트라

https://www.acmicpc.net/problem/1167

1167번: 트리의 지름

문제 상황은 되게 단순하다.

 

가중치가 존재하는 트리 구조가 주어졌을 때,

 

임의의 A 노드부터 B 노드까지의 경로가 존재할 것이고

 

그 중 최장 경로를 트리의 지름이라고 정의한다.

 

트리의 지름을 반환하면 된다.

 

음.. 어떻게 풀지?

 

일단 시간 제한은 2초로 평범한 편이고,

 

데이터를 보니까 노드의 개수는 100000 개이다.

 

에지의 개수는 주어지지 않았으나, 최대 100000 - 1 개가 존재할 수 있을 것이다.

 

어쨋든 간에 최장경로를 구해야 한다고 치면

 

다익스트라를 조금 변형해서 구할 수 있을 텐데

 

다익스트라 알고리즘의 시간 복잡도는 O(VlogE) 이고

 

100000 * Log(100000) ~= 500000 이다.

 

그럼 아무 노드에 대해서 다익스트라 알고리즘을 통해 최장거리를 구하자.

 

그리고 나서 가장 먼 거리에 해당하는 노드를 찾자.

 

그 노드에 대해서 최장거리 다익스트라 알고리즘을 한번 더 수행하면

 

트리의 지름을 구할 수 있다.

 

🧐 왜 두 번 수행해야 할까?

▶ 최장거리를 찾는 함수를 getLongestPath(int v) 라고 하자. 운이 좋게도 v 가 루트 혹은 리프 노드라면 한 번만 수행해도 트리의 지름을 구할 수 있을 것이다. 하지만 어떤 노드가 루트 , 리프 노드인 지는 주어지지 않기 때문에 v 값으로는 임의의 노드를 택할 수 밖에 없다. 임의로 선택한 v 노드가 트리의 한 가운데 위치한 노드라고 상상해보자. 그럼 v 노드에서 다익스트라 알고리즘을 수행하더라도, 트리의 지름을 구할 수 없다. 그러므로 v 노드에서 가장 멀리 있는 노드(A)를 찾고, 해당 노드에서 getLongestPath 를 다시 수행해서 가장 멀리 있는 노드까지의 거리가 해당 트리에서 최장 경로가 될 것이고 이것이 바로 트리의 지름이다. (조금 더 엄밀하게 증명을 하고 싶은데... 지금은 잘 모르겠다 😢) 

 

이렇게 다익스트라 알고리즘을 2 번 수행해서 구할 수도 있고

 

그냥 DFS 로 dist 배열을 갱신해가면서 해결할 수도 있다.

 

DFS 를 수행할 경우 O(n+m) ~= O(200000) 이므로

 

다익스트라 알고리즘 보다는 조금 더 빠른 시간 복잡도로 해결할 수 있을 것이다.

 

다익스트라 알고리즘을 변형하는 방법은 조금 구현이 번거로울 수 있다.

 

그러므로 먼저 DFS 를 이용하여 코드를 작성하였다.

 

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

public class Main {
    static int v;
    static ArrayList<Node>[] graph;
    static int[] dist;
    static boolean[] visited;

    public static void main(String[] args) throws IOException {
        System.setIn(new FileInputStream("src/input.txt"));
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer stk;

        v = Integer.parseInt(br.readLine()); // 정점의 개수
        dist = new int[v + 1];
        visited = new boolean[v + 1];

        graph = new ArrayList[v + 1];
        for (int i = 0; i <= v; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int i = 1; i <= v; i++) {
            stk = new StringTokenizer(br.readLine());

            int a = Integer.parseInt(stk.nextToken());
            while (true) {
                int b = Integer.parseInt(stk.nextToken());
                if (b == -1) break;
                int w = Integer.parseInt(stk.nextToken());
                graph[a].add(new Node(b, w));
            }
        }

        dfs(1);
        int max = -1;
        int max_node = -1;
        for (int i = 1; i <= v; i++) {
            if (dist[i] >= max) {
                max = dist[i];
                max_node = i;
            }
        }

        dist = new int[v + 1];
        visited = new boolean[v + 1];
        dfs(max_node);
        for (int i = 1; i <= v; i++) {
            max = Math.max(max, dist[i]);
        }

        bw.write(max + "\n");
        bw.flush();
        br.close();
        bw.close();
    }

    static class Node {
        int to;
        int w;

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

    static void dfs(int v) {
        visited[v] = true;
        for (Node ele : graph[v]) {
            // 연결된 노드를 탐색
            if (!visited[ele.to]) {
                // 방문하지 않았다면
                visited[ele.to] = true; // 방문
                dist[ele.to] = dist[v] + ele.w;
                dfs(ele.to);
            }
        }
    }
}

위 코드를 이용하면 AC 를 받을 수 있다.

 

 

그러면 이제 다익스트라 알고리즘으로 문제를 해결해보자.

 

dfs 함수 대신에 Dijkstra 함수를 구현해서 대체해주면 된다.

 

주의할 것은 최장거리 다익스트라 알고리즘 이므로

 

초기화 하는 부분이라던가,

 

경로를 탐색하는 부분의 조건문,

 

Node 의 Comparable 을 구현하는 데에서

 

일반적인 최단거리 Dijkstra 알고리즘과 정 반대로 해주어야 한다는 것이다!

 

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

public class Main {
    static int v;
    static ArrayList<Node>[] graph;
    static int[] dist;
    static boolean[] visited;

    public static void main(String[] args) throws IOException {
        System.setIn(new FileInputStream("src/input.txt"));
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer stk;

        v = Integer.parseInt(br.readLine()); // 정점의 개수
        dist = new int[v + 1];
        visited = new boolean[v + 1];

        graph = new ArrayList[v + 1];
        for (int i = 0; i <= v; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int i = 1; i <= v; i++) {
            stk = new StringTokenizer(br.readLine());

            int a = Integer.parseInt(stk.nextToken());
            while (true) {
                int b = Integer.parseInt(stk.nextToken());
                if (b == -1) break;
                int w = Integer.parseInt(stk.nextToken());
                graph[a].add(new Node(b, w));
            }
        }

        dijkstra(new Node(1, 0));
        int max = -1;
        int max_node = -1;
        for (int i = 1; i <= v; i++) {
            if (dist[i] >= max) {
                max = dist[i];
                max_node = i;
            }
        }

        dist = new int[v + 1];
        visited = new boolean[v + 1];

        dijkstra(new Node(max_node, 0));
        for (int i = 1; i <= v; i++) {
            max = Math.max(max, dist[i]);
        }

        bw.write(max + "\n");
        bw.flush();
        br.close();
        bw.close();
    }

    static class Node implements Comparable<Node> {
        int to;
        int w;

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

        @Override
        public int compareTo(Node o) {
            if(this.w <= o.w){
                return 1;
            } return -1;
        }
    }

    static void dijkstra(Node start){
        PriorityQueue<Node> pq = new PriorityQueue<>();
        Arrays.fill(dist, -Integer.MAX_VALUE);
        dist[start.to] = 0;
        visited[start.to] = true;
        pq.add(start);

        while(!pq.isEmpty()){
            Node now = pq.poll();
            for(Node ele : graph[now.to]){
                // 연결된 노드 탐색
                if(!visited[ele.to]){
                    if(dist[ele.to] < dist[now.to] + ele.w){
                        // 더 최장 경로를 발견한다면
                        dist[ele.to] = dist[now.to] + ele.w;
                        visited[ele.to] = true;
                        pq.add(new Node(ele.to, dist[ele.to]));
                    }
                }
            }
        }
    }
}

 

다익스트라 알고리즘을 변형하는 과정에서 실수가 나올 수도 있어서

 

그냥 DFS 로 구현하는 것이 조금 더 간단하기도 하고

 

시간 제한도 조금 더 널널하게 통과할 수 있으므로 좋을 것 같다.

 

2023. 03. 20. 월.