[실전 문제] 트리에서의 거리

1. 문제

트리가 주어지고, 두 노드 X, Y가 주어질 때, 이 두 노드 사이의 거리를 출력하는 프로그램을 작성하시오. 트리에서는 두 노드를 잇는 경로가 유일하기 때문에, 정답은 항상 유일하다는 것을 참고한다. 예를 들어, 다음과 같은 트리에서 노드 3, 노드 6 사이의 거리는 4이다.

---

입력

첫 번째 줄에 트리의 노드 개수 n, 두 노드 X, Y의 번호가 주어진다. ( 0 ≤ X, Y ≤ n < 1000 ) 두 번째 줄부터 트리의 간선 정보가 주어진다. 각 줄은 2개의 숫자 a, b로 이루어지며, 이는 노드 a가 노드 b의 부모노드라는 것을 의미한다. 루트는 노드 0이라고 가정한다. 

출력

두 노드 X, Y 사이의 거리를 출력한다.

예제 입력

11 3 6
0 1
0 2
1 3
1 4
1 5
2 6
2 10
6 7
6 8
6 9

예제 출력

4

 

 

 

2. 풀이

이 문제는 공통 조상 찾기 문제를 응용하면 쉽게 풀 수 있다. 우선 x에서 루트까지 방문하며 방문한 지점에 0부터 값을 증가시키면서 넣어준다. y에서도 마찬가지로 루트까지 방문하면서 값을 넣어주되, x에서 이미 방문한 지점을 만났을 때 계산을 해주면 된다.

 

x와 y의 거리는 반드시 두 노드의 공통된 부모를 거치기 때문에 두 노드의 공통된 부모 노드를 찾는 것이 핵심이다.

import java.io.*;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {

    private static final BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    private static final BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

    private static int n;
    private static int x;
    private static int y;

    private static int[] parent;
    private static int[] visitedX;
    private static int[] visitedY;

    public static void main(String[] args) throws IOException {

        StringTokenizer st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        x = Integer.parseInt(st.nextToken());
        y = Integer.parseInt(st.nextToken());

        parent = new int[n];
        visitedX = new int[n];
        visitedY = new int[n];

        Arrays.fill(visitedX, -1);
        Arrays.fill(visitedY, -1);
        for(int i = 0; i < n-1; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            parent[b] = a; // a 노드가 b 노드의 부모 노드
        }

        int count = 0;
        while(true) {
            visitedX[x] = count++;
            if(x == 0) break;
            x = parent[x];
        }

        count = 0;
        int sum = 0;
        while(true) {
            if(visitedX[y] != -1) {
                sum = visitedX[y] + count;
                break;
            } else {
                visitedY[y] = count++;
                if(y == 0) break;
                y = parent[y];
            }
        }
        
        bw.write(sum + "");
        br.close();
        bw.flush();
        bw.close();
    }

}

---

[참고] 코딩 테스트 - 필수 문제 - 공통 조상 찾기