[PS] 백준 2250번 트리의 높이와 너비

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

문제 설명

이진트리를 다음의 규칙에 따라 행과 열에 번호가 붙어있는 격자 모양의 틀 속에 그렸을 때, 각 레벨의 너비 중 최대 너비를 갖는 레벨과 그 너비를 구하는 문제이다.

1. 이진트리에서 같은 레벨에 있는 노드는 같은 행에 위치한다.
2. 한 열에는 한 노드만 존재한다.
3. 임의의 노드의 왼쪽 서브트리에 있는 노드들은 해당 노드보다 왼쪽의 열에 위치하고, 오른쪽 서브트리에 있는 노드들은 해당 노드보다 오른쪽의 열에 위치한다.
4. 노드가 배치된 가장 왼쪽 열과 오른쪽 열 사이엔 아무 노드도 없이 비어있는 열은 없다.

문제 풀이

먼저 각 노드의 열 번호를 알아야 한다. 문제에서 주어진 그림을 잘 보면, 열 번호는 이진트리를 중위 순회하는 순서와 동일하다는 것을 눈치 챌 수 있다. 따라서, 이진 트리를 중위 순회하여 각 노드의 열 번호를 구한다.

열 번호를 모두 구했다면 이제 각 레벨에서의 너비를 구해야 한다. 레벨 별로 이진트리를 탐색하고 싶으니 BFS를 이용하면서 Queue의 사이즈만큼만 꺼내서 자식을 삽입하는 것을 반복하면 된다. 그리고 Queue에 삽입할 때 왼쪽 자식부터 넣는다면 그 순서대로 값이 나오기 때문에 같은 레벨의 정점의 열 번호를 Deque에 넣고 Deque의 처음과 끝에 위치하는 열 번호를 빼서 1을 더해주면 그 레벨의 너비가 된다.

마지막으로 그 중 넓 큰 너비를 가지는 레벨과 그 너비를 출력하면 된다. 가장 넓은 너비를 가지는 레벨이 다수인 경우, 가장 작은 레벨을 출력한다.

여담

이 문제의 입력에서 루트노드에 대한 언급이 없다. 즉, 1번 노드가 루트노드가 아닐 수 있다는 말이다. 이 포인트 때문에 삽질을 좀 했다. 따라서 중위 순회시와 BFS 수행시에 루트노드가 시작 노드가 되어야 하므로 그 이전에 루트노드를 구하는 작업을 수행해주면 된다.

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.Queue;
import java.util.Scanner;
import java.util.StringTokenizer;

class Main {
	static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
	static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
	static Scanner sc = new Scanner(System.in);
	static StringTokenizer st;

	static class Node {
		int data;
		int col;
		int left;
		int right;

		public Node(int data, int left, int right) {
			super();
			this.data = data;
			this.left = left;
			this.right = right;
		}

	}

	static int n, maxLevel, maxWidth;
	static Node[] nodes;
	static int[] parent;
	static int col = 1;

	public static void main(String[] args) throws Exception {
		n = Integer.parseInt(br.readLine());

		nodes = new Node[n + 1];
		parent = new int[n + 1];

		for (int i = 1; i <= n; i++) {
			st = new StringTokenizer(br.readLine());

			int num = Integer.parseInt(st.nextToken());
			int left = Integer.parseInt(st.nextToken());
			int right = Integer.parseInt(st.nextToken());

			nodes[num] = new Node(num, left, right);

			if (left != -1) {
				parent[left] = num;
			}

			if (right != -1) {
				parent[right] = num;
			}
		}

		int rootNum = findRoot();
		inOrder(rootNum);
		go(rootNum);

		System.out.println(maxLevel + " " + maxWidth);
	}

	/**
	 * 임의의 정점에서 루트노드일 때까지 타고 올라감
	 * 단, 입력 정점 번호가 1~N 이므로 항상 존재하는 1번 정점에서 시작
	 * 
	 * @return 루트노드의 번호
	 */
	static int findRoot() {
		int curNum = 1;
		while (true) {
			if (parent[curNum] == 0) {
				return curNum;
			}

			curNum = parent[curNum];
		}
	}

	/**
	 * 중위순회로 노드의 열 번호를 구함
	 * 
	 * @param idx
	 */
	static void inOrder(int idx) {
		if (idx == -1) {
			return;
		}

		inOrder(nodes[idx].left);
		nodes[idx].col = col++;
		inOrder(nodes[idx].right);
	}

	/**
	 * BFS로 레벨 별 탐색 트리이므로 방문체크 불필요
	 */
	static void go(int rootNum) {
		Queue<Integer> queue = new ArrayDeque<>();
		queue.add(rootNum);

		int level = 1;
		while (!queue.isEmpty()) {
			int size = queue.size(); // 같은 레벨인 정점의 개수를 구함

			Deque<Integer> deque = new ArrayDeque<>(); // 같은 레벨인 모든 정점들의 열 번호를 저장
			for (int i = 0; i < size; i++) { // // 같은 레벨인 정점의 개수만큼 꺼냄
				int cur = queue.poll();

				deque.add(nodes[cur].col);

				int left = nodes[cur].left;
				if (left != -1) {
					queue.add(left);
				}

				int right = nodes[cur].right;
				if (right != -1) {
					queue.add(right);
				}
			}

			int width = deque.getLast() - deque.getFirst() + 1; // 왼쪽 자식부터 Queue에 삽입하므로 꺼낼 때도 왼쪽 정점부터 꺼냄(최우측 - 최좌측 + 1)
			if (width > maxWidth) {
				maxWidth = width;
				maxLevel = level;
			}

			level++;
		}

	}
}