[백준] 1167 - 트리의 지름 (자바/Java)

 

문제 링크

성능 요약

메모리: 87924 KB, 시간: 820 ms

분류

깊이 우선 탐색(dfs), 그래프 이론(graphs), 그래프 탐색(graph_traversal), 트리(trees)

문제 설명

트리의 지름이란, 트리에서 임의의 두 점 사이의 거리 중 가장 긴 것을 말한다. 트리의 지름을 구하는 프로그램을 작성하시오.

입력

트리가 입력으로 주어진다. 먼저 첫 번째 줄에서는 트리의 정점의 개수 V가 주어지고 (2 ≤ V ≤ 100,000)둘째 줄부터 V개의 줄에 걸쳐 간선의 정보가 다음과 같이 주어진다. 정점 번호는 1부터 V까지 매겨져 있다.

먼저 정점 번호가 주어지고, 이어서 연결된 간선의 정보를 의미하는 정수가 두 개씩 주어지는데, 하나는 정점번호, 다른 하나는 그 정점까지의 거리이다. 예를 들어 네 번째 줄의 경우 정점 3은 정점 1과 거리가 2인 간선으로 연결되어 있고, 정점 4와는 거리가 3인 간선으로 연결되어 있는 것을 보여준다. 각 줄의 마지막에는 -1이 입력으로 주어진다. 주어지는 거리는 모두 10,000 이하의 자연수이다.

출력

첫째 줄에 트리의 지름을 출력한다.

---

많은 시행착오를 겪었다..

처음에는 dfs를 1번만 하면 될 줄 알았다. 그런데 주어진 테스트 케이스 외에 dfs 시작 노드인 1이 루트인 경우에는 최대 거리를 반영하지 못하는 문제가 있었다.

그래서 dfs를 리프 노드부터 수행하면 되지 않을까 했는데, 모든 리프노드에 대해 dfs를 수행하면 시간 초과이고, 리프 노드를 하나만 골라서 하면 그 노드에서 시작하는 지름이 최대 거리가 아닐 경우도 존재했다.

그럼 트리를 구현해야 하나 생각했는데, 주어지는 것은 edge들이고 루트 노드를 알 수 없기 때문에 트리를 구현하는 것도 어려워 보였다.

 

그래서 결국 다른 풀이를 찾아보았는데, dfs를 두 번 하면 해결되는 간단한 문제였다.

특정 노드에서 dfs를 진행한 후 최대 지름인 노드를 찾는다. 그 노드에서 다시 dfs를 하면 최대 거리를 찾을 수 있었다.

트리는 구현하는 방식에 따라 어떤 노드든 root 노드가 될 수 있다. 그러므로 한 노드에서 dfs를 하며 찾아낸 지름이 가장 긴 노드는, root 노드로 삼아 처음 dfs를 시작한 노드에서 한 쪽 방향으로 가장 먼 노드를 의미한다.

거기서 다시 dfs를 한다면 기존의 root 노드를 거쳐 다른 쪽의 먼 노드까지 도달할 수 있으므로 이때 구한 값이 지름에 해당한다.

여러가지 방안을 생각했지만, dfs를 두 번 한다는 건 생각하기 어려웠던 것 같다

 

• 입력

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		N = Integer.parseInt(br.readLine());
		nodes = new TreeNode[N + 1];
		visited = new boolean[N + 1];
		for (int i = 1; i <= N; i++) {
			nodes[i] = new TreeNode(i);
		}
		for (int i = 0; i < N; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			int count = (st.countTokens() - 1) / 2;
			int start = Integer.parseInt(st.nextToken());
			int end = Integer.parseInt(st.nextToken());
			nodes[start].childs = new TreeNode[count];
			nodes[start].weights = new int[count];
			int idx = 0;
			while (true) {
				if (end == -1)
					break;
				int weight = Integer.parseInt(st.nextToken());
				nodes[start].childs[idx] = nodes[end];
				nodes[start].weights[idx++] = weight;
				end = Integer.parseInt(st.nextToken());
			}
		}

nodes라는 배열에 TreeNode들을 담아주었다. 1번부터 N번까지 존재하기 때문에 각 index는 그 노드의 data가 된다.

그리고 입력을 돌면서 count 변수를 통해 해당 start 노드에 연결된 edge의 개수를 세어 배열을 만들었다.

childs는 edge에 연결된 다른 노드를, weight는 그 edge의 거리(가중치)를 의미한다.

 

TreeNode 클래스는 다음과 같이 data, 연결된 노드를 담을 childs, 그 노드와의 거리를 담을 weights 배열로 구성되었다.

class TreeNode {
	int data;
	TreeNode[] childs;
	int[] weights;

	public TreeNode(int data) {
		this.data = data;
	}

}

dfs는 다음과 같이 구현하였다. 먼저 해당 idx의 TreeNode를 찾는다.

idx가 0일 때는, N이 1일 때, 노드가 하나여서 other(dfs의 시작 노드(1)에서 가장 먼 노드)가 0일 경우, node의 childs가 존재하지 않아 NullPointer Exception이 나올 것을 방지하기 위해 처리한 것이다.

 

node의 childs 배열을 통해 edge들을 파악하고, 방문되지 않았을 경우 거리를 더해주며 ans와 비교한다.

이때 거리가 ans보다 클 경우에는 other와 ans를 업데이트한다.

	public static void dfs(int idx, int sum) {
		TreeNode node = nodes[idx];
		visited[idx] = true;

		if (idx == 0)
			return; // N이 1일 때 처리

		for (int i = 0; i < node.childs.length; i++) {
			TreeNode child = node.childs[i];
			if (!visited[child.data]) {
				visited[child.data] = true;
				if (sum + node.weights[i] > ans) {
					ans = sum + node.weights[i]; // 길이 update
					other = child.data; // other node update
				}
				dfs(child.data, sum + node.weights[i]);
				visited[child.data] = false;
			}
		}
	}

전체 코드

package Gold.g2;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class BOJ_1167_트리지름 {
	static int N;
	static TreeNode[] nodes;
	static boolean[] visited;

	static int ans;
	static int other = 0;

	public static void main(String[] args) throws NumberFormatException, IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		N = Integer.parseInt(br.readLine());
		nodes = new TreeNode[N + 1];
		visited = new boolean[N + 1];
		for (int i = 1; i <= N; i++) {
			nodes[i] = new TreeNode(i);
		}
		for (int i = 0; i < N; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			int count = (st.countTokens() - 1) / 2;
			int start = Integer.parseInt(st.nextToken());
			int end = Integer.parseInt(st.nextToken());
			nodes[start].childs = new TreeNode[count];
			nodes[start].weights = new int[count];
			int idx = 0;
			while (true) {
				if (end == -1)
					break;
				int weight = Integer.parseInt(st.nextToken());
				nodes[start].childs[idx] = nodes[end];
				nodes[start].weights[idx++] = weight;
				end = Integer.parseInt(st.nextToken());
			}
		}
		dfs(1, 0);
		Arrays.fill(visited, false);
		dfs(other, 0); // 루트노드에서 가장 먼 노드부터 다시 dfs 시행
		System.out.println(ans);
	}

	public static void dfs(int idx, int sum) {
		TreeNode node = nodes[idx];
		visited[idx] = true;

		if (idx == 0)
			return; // N이 1일 때 처리

		for (int i = 0; i < node.childs.length; i++) {
			TreeNode child = node.childs[i];
			if (!visited[child.data]) {
				visited[child.data] = true;
				if (sum + node.weights[i] > ans) {
					ans = sum + node.weights[i]; // 길이 update
					other = child.data; // other node update
				}
				dfs(child.data, sum + node.weights[i]);
				visited[child.data] = false;
			}
		}
	}
}

class TreeNode {
	int data;
	TreeNode[] childs;
	int[] weights;

	public TreeNode(int data) {
		this.data = data;
	}

}