본문 바로가기

코딩테스트/JAVA

[백준 JAVA] 1167번: 트리의 지름 (두가지 풀이 방법 1967번과 비교)

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

 

1167번을 풀면서 처음에는 예전에 풀었던 1967번 트리의 지름 풀이를 그대로 가져와도 되지 않을까 생각했다. 실제로 아이디어 자체는 꽤 비슷하다. 각 정점에서 아래로 내려가는 가장 긴 경로와 두 번째로 긴 경로를 구하고, 그 둘의 합으로 지름 후보를 갱신하는 방식이다.

 

1967번: 트리의 지름 풀이 링크

https://dwshin-dev.tistory.com/257

 

그런데 1967에서는 잘 되던 코드가 1167에서는 바로 통하지 않았다. 처음에는 단순히 방문 배열을 안 써서 그런가 싶었고, 더 생각해보니 그 차이가 맞긴 했지만 실제로는 입력으로 주어지는 트리의 형태 자체가 달랐다. 이 글에서는 그 과정에서 어떻게 생각이 바뀌었는지, 그리고 마지막에 왜 정석 풀이로 방향을 틀었는지를 정리해보려고 한다.

 

문제 요약

트리의 모든 정점과 가중치가 있는 간선 정보가 주어질 때, 트리의 지름을 구하는 문제다. 트리의 지름은 임의의 두 정점 사이의 거리 중 가장 긴 값을 의미한다.

 

처음 떠올린 방법

예전에 1967번에서 사용했던 방식은 이런 흐름이었다.

  • 어떤 정점 u에서 자식 방향으로 내려간다.
  • u를 루트로 하는 서브트리에서 가장 긴 경로 길이를 반환한다.
  • 자식들로부터 올라온 길이 중 가장 큰 값 first와 두 번째 큰 값 second를 관리한다.
  • first + second가 그 정점을 지나가는 가장 긴 경로이므로 전체 최대값을 갱신한다.

1967에서는 이 방식이 자연스럽게 통했다. 실제 코드도 꽤 단순했다.

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

public class Main {
    static int N, max;
    static List<int[]>[] tree;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        N = Integer.parseInt(br.readLine());
        tree = new List[N + 1];
        for (int i = 1; i <= N; i++) tree[i] = new ArrayList<>();

        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(br.readLine());
            int p = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());
            int w = Integer.parseInt(st.nextToken());
            tree[p].add(new int[] {c, w});
        }

        max = 0;
        maxLength(1);

        System.out.println(max);
    }

    private static int maxLength(int u) {
        int first = 0;
        int second = 0;

        for (int[] next : tree[u]) {
            int v = next[0];
            int w = next[1];

            int length = maxLength(v) + w;

            if (length > first) {
                second = first;
                first = length;
            } else if (length > second) {
                second = length;
            }
        }

        max = Math.max(max, first + second);
        return first;
    }
}

 

이 방식이 1967에서 됐던 이유

핵심은 입력이 사실상 부모와 자식 관계로 주어진다는 점이다.

1967에서는 간선을 부모 -> 자식 방향으로만 저장해도 된다. 그러면 DFS는 자연스럽게 아래 방향으로만 진행되고, 다시 부모로 되돌아갈 일이 없다. 그래서 방문 배열이 없어도 문제없이 동작한다.

즉, 1967에서는 트리를 루트가 있는 구조처럼 다뤄도 괜찮았다.

 

1167에서 같은 방식이 바로 안 되는 이유

1167은 입력 형식이 다르다. 각 정점마다 연결된 모든 인접 정점 정보가 들어온다. 즉, 부모와 자식이 구분된 트리가 아니라 무방향 트리의 인접 리스트가 주어진다.

예를 들어 1번 정점에 3번이 연결되어 있다면, 3번 정점 쪽에서도 다시 1번이 보인다. 이 상태에서 1967처럼 아무 처리 없이 재귀를 돌리면 자식으로 내려갔다가 다시 부모로 올라가고, 다시 내려가는 식으로 꼬이게 된다.

 

그래서 처음에는 "결국 visited 하나 쓰면 되는 거 아닌가"라는 생각이 들었는데 사실 그 말은 맞다. 개념적으로 가장 큰 차이는 그 부분이다. 다만 실전에서는 그 차이가 꽤 크다. 1967은 애초에 되돌아가는 간선을 신경 쓸 필요가 없었고, 1167은 반드시 부모로 역주행하는 경로를 막아야 한다.

 

1167도 같은 아이디어로 풀 수는 있다

즉, 1967에서 썼던 first, second 방식이 틀린 건 아니다. 1167에서도 부모 정점을 인자로 넘겨서 되돌아가는 간선만 막아주면 그대로 사용할 수 있다.

이 방식으로 고치면 다음과 같다.

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

public class Main {
    static int V, max;
    static List<int[]>[] tree;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        V = Integer.parseInt(br.readLine());
        tree = new List[V + 1];
        for (int i = 1; i <= V; i++) tree[i] = new ArrayList<>();

        for (int i = 0; i < V; i++) {
            st = new StringTokenizer(br.readLine());
            int cur = Integer.parseInt(st.nextToken());

            while (true) {
                int x = Integer.parseInt(st.nextToken());
                if (x == -1) break;
                int w = Integer.parseInt(st.nextToken());
                tree[cur].add(new int[] {x, w});
            }
        }

        maxLength(1, 0);
        System.out.println(max);
    }

    private static int maxLength(int u, int parent) {
        int first = 0;
        int second = 0;

        for (int[] next : tree[u]) {
            int v = next[0];
            int w = next[1];

            if (v == parent) continue;

            int length = maxLength(v, u) + w;

            if (length > first) {
                second = first;
                first = length;
            } else if (length > second) {
                second = length;
            }
        }

        max = Math.max(max, first + second);
        return first;
    }
}

1967에서 쓰던 아이디어를 거의 그대로 유지하면서, 부모 정점만 제외하는 방식으로 바꾼 것이다.

그런데 여기서 한 번 더 고민하게 된다. 굳이 이렇게 풀어야 할까?

 

왜 정석 풀이를 따로 고려하게 됐는가

1167은 정점 수가 많고 입력도 무방향 트리 형태라서 구현할 때 조금만 삐끗해도 방문 처리 시점이나 부모 제외 처리에서 실수가 생기기 쉽다. 실제로 처음 풀이를 짤 때도 이 부분에서 많이 헷갈렸다.

 

또 자바에서는 재귀 깊이가 깊어지면 스택 오버플로우 가능성도 신경 쓰게 된다. 그래서 이 문제는 트리 DP 느낌으로 푸는 것도 가능하지만 많은 사람들이 더 단순한 정석 풀이를 선호한다.

 

그 정석 풀이가 바로 "가장 먼 정점 찾기 두 번"이다.

정석 풀이의 아이디어

이 방식은 다음 두 단계로 끝난다.

  1. 아무 정점 하나에서 시작해서 가장 먼 정점 A를 찾는다.
  2. A에서 다시 가장 먼 정점까지의 거리를 구한다.
  3. 그 거리가 트리의 지름이다.

처음 보면 왜 이게 되는지 조금 낯설 수 있는데, 트리에서는 임의의 한 정점에서 가장 멀리 있는 정점이 지름의 한 끝점이 되는 성질이 있다. 그래서 한 번 가장 먼 점을 찾고 거기서 다시 한 번 가장 먼 점까지 가면 지름이 된다.

 

이 방식의 장점은 생각할 게 적다는 점이다. first, second를 관리할 필요도 없고, 어떤 정점이 루트인지도 중요하지 않다. 무방향 트리에서는 오히려 이 풀이가 훨씬 자연스럽다.

 

DFS 재귀 를 이용한 정석 풀이

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

public class Main {
    static int V;
    static List<int[]>[] tree;
    static int[] dist;
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        V = Integer.parseInt(br.readLine());
        tree = new List[V + 1];
        for (int i = 1; i <= V; i++) tree[i] = new ArrayList<>();

        for (int i = 0; i < V; i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            while(true) {
                int v = Integer.parseInt(st.nextToken());
                if (v == -1) break;
                int w = Integer.parseInt(st.nextToken());
                tree[u].add(new int[] {v, w});
            }
        }

        int[] first = findFarNode(1);
        int[] second = findFarNode(first[0]);

        System.out.println(second[1]);
    }

    private static int[] findFarNode(int start) {
        dist = new int[V + 1];
        Arrays.fill(dist, -1);

        dist[start] = 0;
        dfs(start);

        int farNode = start;
        for (int i = 1; i <= V; i++) {
            if (dist[i] > dist[farNode]) {
                farNode = i;
            }
        }

        return new int[] {farNode, dist[farNode]};
    }

    private static void dfs(int u) {
        for (int[] next : tree[u]) {
            int v = next[0];
            int w = next[1];

            if (dist[v] != -1) continue;

            dist[v] = dist[u] + w;
            dfs(v);
        }
    }

}

 

재귀 대신 스택을 이용한 반복 DFS로 구현하면 자바에서도 더 안전하다.

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

public class Main {
    static int V;
    static List<int[]>[] tree;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        V = Integer.parseInt(br.readLine());
        tree = new List[V + 1];
        for (int i = 1; i <= V; i++) tree[i] = new ArrayList<>();

        for (int i = 0; i < V; i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            while(true) {
                int v = Integer.parseInt(st.nextToken());
                if (v == -1) break;
                int w = Integer.parseInt(st.nextToken());
                tree[u].add(new int[] {v, w});
            }
        }

        int[] first = findFarNode(1);
        int[] second = findFarNode(first[0]);

        System.out.println(second[1]);
    }

    private static int[] findFarNode(int start) {
        int[] dist = new int[V + 1];
        Arrays.fill(dist, -1);

        Deque<Integer> stack = new ArrayDeque<>();
        stack.push(start);
        dist[start] = 0;

        while (!stack.isEmpty()) {
            int u = stack.pop();

            for (int[] next : tree[u]) {
                int v = next[0];
                int w = next[1];

                if (dist[v] != -1) continue;

                dist[v] = dist[u] + w;
                stack.push(v);
            }
        }

        int farNode = start;
        for (int i = 1; i <= V; i++) {
            if (dist[i] > dist[farNode]) {
                farNode = i;
            }
        }

        return new int[] {farNode, dist[farNode]};
    }
}

 

이 풀이가 더 마음에 들었던 이유

처음에는 1967 방식과 1167 방식이 본질적으로 크게 다르지 않다고 생각했다. 실제로도 그렇다. 1167 역시 부모만 잘 막아주면 같은 논리로 풀 수 있다.

하지만 문제를 풀다 보면 정답을 아는 것과 실전에서 덜 헷갈리게 구현하는 것은 조금 다르다. 1167은 입력이 무방향 트리라서, 어떤 방향으로 내려가야 한다는 보장이 처음부터 없다. 그래서 트리 DP처럼 아래로 내려가는 구조를 인위적으로 만들어야 한다.

반면 두 번 탐색하는 방식은 그런 고민이 없다. 그냥 연결된 정점을 따라 가장 먼 곳까지 가고, 거기서 한 번 더 가면 끝이다. 구현도 단순하고, 실수할 여지도 적다.

시간복잡도

두 방식 모두 시간복잡도는 O(V)다.

  • first, second를 관리하는 방식도 모든 정점을 한 번씩 방문한다.
  • 가장 먼 정점 찾기 두 번 방식도 탐색을 두 번 할 뿐이라 전체는 O(V)다.

따라서 성능 차이보다는 구현 안정성과 실수 가능성 차이로 선택하는 편이 좋다.

 

정리

이 문제를 풀면서 느낀 건, 예전에 맞았던 풀이를 그대로 다른 문제에 적용하는 습관이 꼭 나쁜 것은 아니지만, 입력 구조가 달라졌다면 그 차이를 먼저 봐야 한다는 점이다.

1967에서는 트리를 루트가 있는 형태처럼 다뤄도 문제가 없었지만, 1167은 무방향 트리라서 부모로 되돌아가는 경로를 반드시 막아야 했다. 그래서 겉보기에는 비슷한 지름 문제라도 구현 방식은 꽤 달라질 수 있다.

결국 1167도 1967 방식으로 풀 수는 있지만, 정석은 두 번 탐색하는 방식이고 실제로 구현도 더 깔끔하다. 이번 문제는 틀린 아이디어를 잡았다기보다는, 같은 개념을 어떤 문제에서는 어떻게 구현해야 하는지를 다시 한 번 배우게 된 문제였다.

 

시간초는 밑에서 부터 각각 기존 1967번 풀이를 이용한 방법, 재귀를 이용한 dfs, 스택을 이용한 dfs이다.

 

 

시간은 스택을 사용한게 가장 오래걸렸는데 Integer 오토박싱/언박싱 + ArrayDeque 조작 비용때문이라는 생각이 들었다.

따라서 풀이할때 적절히 안정적인 방법을 선택할것인지 시간을 중요시 생각할것인지 판단해서 풀이 방법을 선택하면 좋을꺼같다.

만약 안정성과 시간 둘다 동시에 챙기고 싶다면 int 배열 형태의 스택을 구현해서 사용하는 방법도 있을꺼같다.

 

회고

처음에는 "visited 하나만 추가하면 되는 거 아닌가"라고 생각했고, 실제로 절반은 맞는 말이었다. 다만 코딩테스트에서는 맞는 아이디어보다 덜 헷갈리는 구현이 더 중요할 때가 많다. 이번 문제는 그걸 잘 보여줬다. 다음에 비슷한 트리 지름 문제가 나오면 무작정 예전 풀이를 꺼내기보다, 입력이 rooted tree인지 undirected tree인지부터 먼저 확인하고 시작해야겠다.