[DFS] 백준 - 1967번: 트리의 지름 Java 풀이

[공부용이기 때문에 코드가 깔끔하지 않을 수 있습니다!]

 

1967번: 트리의 지름

 

 

문제는 다음과 같다.

간선의 길이를 다 합쳤을 때 가장 긴 지름인 두 노드를 구해서 그 지름을 프린트하는 것이다!

DFS로 금방 풀 수 있다.

 

코드

import java.util.ArrayList;
import java.util.Scanner;

public class Main {
    static ArrayList<Node>[] map;
    static boolean[] visit;
    static int n;
    static int max;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        n = sc.nextInt();

        map = new ArrayList[n];
        visit = new boolean[n];

        for (int i = 0; i < n; i++) {
            map[i] = new ArrayList<>();
        }

        for (int i = 0; i < n - 1; i++) {
            int a = sc.nextInt() - 1;
            int b = sc.nextInt() - 1;
            int d = sc.nextInt();

            map[a].add(new Node(b, d));
            map[b].add(new Node(a, d));
        }

        for (int i = 0; i < n; i++) {
            if (!visit[i])
                dfs(i, 0);
        }

        System.out.println(max);
    }

    private static void dfs(int start, int d) {
        if (d > max) {
            max = d;
            return;
        }
        visit[start] = true;

        for (int i = 0; i < map[start].size(); i++) {
            if (!visit[map[start].get(i).next]) {
                dfs(map[start].get(i).next, d + map[start].get(i).dist);
            }
        }

        visit[start] = false;
    }
}

class Node {
    int next;
    int dist;

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

 

처음에 생각 없이 2차원 배열로 트리를 구성해서 인덱스에 거리를 넣으려고 했다가 메모리가 엄청 많아질 것 같아 노드 클래스를 하나 만들어 리스트 배열을 사용했다.

 

저번에 스태틱을 너무 많이 쓴 것 같아 이번에는 스태틱을 하나라도 줄이기 위해 dfs에서 파라미터 값을 거리 총합을 넣어줬다.

 

그리고 DFS를 돌려주는데 이때 주의할 게 if문으로 종료 지점을 정해줘야 한다.

처음에 종료 지점을 안 넣어주고 visit[start] = false; 위에서 그냥 계속 Math.max(max, d);로 최대값 최신화를 시켜줬는데 이러면 시간 초과된다!

따라서 꼭 if문으로 종료 지점을 정해줘야 한다.