[BOJ] 백준 2132번 나무 위의 벌레 (Java)

#2132 나무 위의 벌레

난이도 : 골드 3

유형 : 트리 / BFS / DFS

2132번: 나무 위의 벌레

▸ 문제

전산학(Computer science)에서 트리란 사이클이 없는 그래프를 말한다. 트리(Tree)라는 이름이 의미하듯, 이러한 구조는 나무의 모습에서 유래한다. 즉, 트리의 각 간선(edge)들이 나무의 가지를 나타내고, 각 정점(node)들은 가지가 갈라지는 지점을 의미한다. 또한 트리의 루트는 나무의 뿌리를 의미한다. 이러한 구조는 일반적인 나무의 구조에 해당하지만, 트리 자체의 성질에 주목하면 실제 나무와는 다소 다른 구조가 되기도 한다.

우리가 생각하려는 나무는 루트가 없는 트리이다. 이때 트리의 각각의 간선은 나무의 가지에 해당하고, 트리의 각 정점은 나무 위에서 열매가 매달려있는 지점을 의미한다. 각각의 정점에는 몇 개의 열매가 매달려 있다. 물론 열매 없이 가지가 갈라지는 경우도 있으므로, 이러한 경우는 그 노드에 0개의 열매가 매달려 있다고 생각하기로 하자.

이러한 나무 위에 한 마리의 벌레가 있다. 이 벌레는 임의의 정점에서 이동하기 시작한다. 벌레가 한 정점에 있을 때에는, 그 정점에 있는 열매들을 먹을 수 있다. 열매들을 다 먹은 후에는 가지를 따라서 다른 정점으로 이동한다. 만약 이동할 수 있는 가지가 여러 개 있다면 그 중 하나를 임의로 선택하지만, 한 번 지났던 가지는 다시 지날 수 없다. 벌레의 이동은 더 이상 이동할 수 있는 정점이 없을 때에 끝난다.

나무의 모양이 주어졌을 때, 벌레가 최대로 먹을 수 있는 열매의 수와 이때 어느 정점에서 이동을 시작해야 하는지를 알아내는 프로그램을 작성하시오.

▸ 입력

첫째 줄에는 트리의 정점의 개수를 나타내는 정수 n(1 ≤ n ≤ 10,000)이 주어진다. 다음 줄에는 차례로 1번, 2번, …, n번 정점에 매달려 있는 열매의 개수가 주어진다. 다음 n-1개의 줄에는 트리의 각 간선을 나타내는 서로 다른 두 자연수 A, B(1 ≤ A, B ≤ n)가 주어진다. 이는 트리의 A번 정점과 B번 정점이 연결되어 있음을 의미한다. 나무에 매달려 있는 열매의 총 개수는 231-1 (2,147,483,647)개를 넘지 않는다.

▸ 출력

첫째 줄에 벌레가 먹을 수 있는 열매의 최대 개수와, 이때 이동을 시작할 정점의 번호를 출력한다. 답이 여러 개 있을 경우에는 정점의 번호가 가장 작은 경우를 출력한다.

 

문제 풀이  

트리의 지름을 찾는 문제이다. 트리의 지름이란 트리에서 거리가 가장 긴 두 정점의 거리를 나타낸다. 다음과 같은 방식에 따르면 dfs 2번만 돌려서 해결할 수 있다.

1. 트리는 어떤 점도 루트가 될 수 있다. 따라서 처음에 아무 정점(u)을 고른다.
2. u에서 가장 거리가 먼 정점(v)을 찾는다.
3. 그러면 u - v 사이의 거리가 트리의 지름이다.

 

이에 대한 귀류법 증명은 해당 글에 가면 있다. 그냥 나는 원을 그려서 기하학적으로 증명을 해봤다. 트리에서 가장 긴 거리를 가지는 두 정점으로 만든 원을 그리면 다음과 같다.

• 가장 긴 거리를 두 정점으로 그린 원이기 때문에 원 밖에는 절대 다른 정점이 존재할 수 없다.

 

가장 긴 거리를 가지는 트리의 두 정점으로 원을 그렸을 때

 

그러면 이제 아무 점이나 잡고(루트) 가장 먼 정점을 찾으면 다음과 같이 2개가 나올 것이다. (만약 다른 지름도 존재하면 복수개 일 수도 있음)

• 만약 주황색 정점을 u로 정한다면 v는 빨간 정점이 될 것이다.
• 만약 초록색 정점을 u로 정한다면 v는 파란 정점이 될 것이다.

 

각 u에 따른 v를 찾는 경우

 

 

그럼 마지막으로 이제 여기서 가장 긴 정점을 찾게되면 결국 파란 정점과 빨간 정점으로 지름이 될 수 밖에 없다.

 

지름 찾기

 

설계

1. 각 정점에 대한 비중과 연결 간선을 저장한다. 
2. 아무 정점(1)을 대상으로 가장 먼 정점을 찾는다.
   1. 여러 개의 정점이 존재할 경우 가장 작은 인덱스를 가진 정점을 구한다.
3. 가장 먼 정점으로 다시 가장 먼 정점(e)를 찾는다.
   1. 2번에서 한 방식과 일치한다.
4. 그렇게 구한 지름 길이 r과 지름 중 가장 작은 정점 번호를 출력한다.

 

풀이 코드 (BFS, DFS)

BFS

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

public class Main {

	static int n;
	static int[] data, vertex;
	static Queue<Integer> q;
	static List<Integer>[] list;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		
		n = Integer.parseInt(br.readLine());
		
		data = new int[n+1];
		vertex = new int[n+1];
		list = new ArrayList[n+1];
		StringTokenizer st = new StringTokenizer(br.readLine());
		for(int i=1; i<n+1; i++) {
			vertex[i] = Integer.parseInt(st.nextToken());
			list[i] = new ArrayList<>();
		}
		
		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());
			
			list[a].add(b);
			list[b].add(a);
		}

		int[] f =  bfs(1);
		int[] r =  bfs(f[1]);
		System.out.println(r[0] +" " + Math.min(f[1], r[1]));
	}
	
	
	static int[] bfs(int s) {
		int idx = 0, maxDis = -1;
		Queue<Integer> q = new LinkedList<>();
		Arrays.fill(data, -1);
		q.add(s);
		data[s] = vertex[s];
		while(!q.isEmpty()) {
			int pos = q.poll();
			
			if(data[pos] > maxDis) {
				idx = pos;
				maxDis = data[pos];
			}else if(data[pos] == maxDis) {
				if(idx > pos) idx = pos;
				maxDis = data[pos];
			}
			
			for(int nxt : list[pos]) {
				if(data[nxt] != -1) continue;
				data[nxt] = data[pos] + vertex[nxt];
				q.add(nxt);
				
			}
		}
		return new int[] {maxDis, idx};
	}
	
}

 

DFS

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

public class Main {

	static int r, s, e;
	static int[] data;
	static List<int[]>[] list;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		
		int n = Integer.parseInt(br.readLine());
		
		data = new int[n+1];
		list = new ArrayList[n+1];
		StringTokenizer st = new StringTokenizer(br.readLine());
		for(int i=1; i<n+1; i++) {
			data[i] = Integer.parseInt(st.nextToken());
			list[i] = new ArrayList<>();
		}
		
		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());
			
			list[a].add(new int[] {b, data[b]});
			list[b].add(new int[] {a, data[a]});
		}
		
		s = 1; e = 10001;
		dfs(1, 0, data[1]);
		
		s = e; r = 0; e = 10001;
		dfs(s, 0, data[s]);
		
		System.out.println(r +" " +  Math.min(s, e));
	}
	
	static void dfs(int idx, int pa, int cost) {
		if(cost > r) {
			r = cost;
			e = idx;
		}else if(cost == r) {
			e = Math.min(e, idx);
		}
		
		for(int[] nxt : list[idx]) {
			if(nxt[0] != pa) {
				dfs(nxt[0], idx, cost+nxt[1]);
			}
		}
	}
}