[BOJ 2957] 이진 탐색 트리 (Java)

BOJ 2957 이진 탐색 트리

문제풀이

가장 먼저 시도한 것은 문제에 나와있는 슈도코드를 그대로 옮겨보았다. 하지만 당연히 시간초과, 최악의 경우에 O(N^2)의 시간복잡도를 가지기 때문이다. 어떻게 접근해야할지 모르겠어서 찾아봤더니 이진 탐색 트리의 삽입 연산에는 규칙이 있었다.

• 삽입하려는 수 x보다 크면서 가장 작은 수와 x보다 작으면서 가장 큰 수 중에서 깊이가 더 깊은 노드의 자식으로 삽입되는 것이다.
• 따라서 직접 자료구조를 구현할 필요 없이 삽입되는 수들의 깊이를 저장하고 있으면 위의 규칙에 따라서 부모가 될 숫자의 깊이 + 1로 삽입하면 된다.

C++의 경우에는 STL에 upper_bound라는 함수가 있어서 x보다 크면서 가장 작은 수를 손쉽게 구할 수 있었다. 그래서 찾아보고 직접 구현해봤다. 그런데 이 함수는 이진탐색을 바탕으로 하기때문에 정렬이 되어있어야한다. 그래서 upper_bound 함수 호출 전에 정렬을 해줬는데 이 때문에 시간초과가 나왔다.

한참을 고민하던 중 자바에도 정렬된 형태로 삽입되는 TreeMap이 있다는 것이 떠올랐다. 그래서 어떤 매서드를 가지고 있는지 확인해봤더니 higerKey(), lowerKey()라는 함수를 가지고 있었다! 바로 적용하였더니 통과!

구현코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.TreeMap;

public class Main {
	
	static TreeMap<Integer, Integer> map;
	static int N;
	static long ans;
	
	public static void main(String[] args) throws NumberFormatException, IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringBuilder sb = new StringBuilder();
		
		N = Integer.parseInt(br.readLine());
		map = new TreeMap<>();
	
		for(int i = 0 ; i < N ; ++i) {
			int x = Integer.parseInt(br.readLine());
			
			if(i == 0) {
				map.put(x, 0);
				ans = 0;
			} else {
				Integer upperKey = map.higherKey(x);
				Integer lowerKey = map.lowerKey(x);
				
				int depth = 0;
				if(upperKey == null) {
					depth = map.get(lowerKey) + 1;
					map.put(x, depth);
				} else if(lowerKey == null) {
					depth = map.get(upperKey) + 1;
					map.put(x, depth);
				} else {
					int upper = map.get(upperKey);
					int lower = map.get(lowerKey);
					
					depth = upper > lower ? upper + 1 : lower + 1;
					map.put(x, depth);
				}
				ans += depth;
			}
			sb.append(ans + "\n");
		}
		
		System.out.println(sb);
	}
}