[BOJ] 백준 11812번 K진 트리 (Java)

#11812 K진 트리

난이도 : 골드 3

유형 : 트리 / LCA

11812번: K진 트리

▸ 문제

각 노드가 자식을 최대 K개 가질 수 있는 트리를 K진 트리라고 한다. 총 N개의 노드로 이루어져 있는 K진 트리가 주어진다.

트리는 "적은 에너지" 방법을 이용해서 만든다. "적은 에너지" 방법이란, 이전 깊이를 모두 채운 경우에만, 새로운 깊이를 만드는 것이고, 이 새로운 깊이의 노드는 가장 왼쪽부터 차례대로 추가 한다.

아래 그림은 노드 9개로 이루어져 있는 3진 트리이다.

노드의 개수 N과 K가 주어졌을 때, 두 노드 x와 y 사이의 거리를 구하는 프로그램을 작성하시오.

▸ 입력

첫째 줄에 N (1 ≤ N ≤ 10^15)과 K (1 ≤ K ≤ 1 000), 그리고 거리를 구해야 하는 노드 쌍의 개수 Q (1 ≤ Q ≤ 100 000)가 주어진다.

다음 Q개 줄에는 거리를 구해야 하는 두 노드 x와 y가 주어진다. (1 ≤ x, y ≤ N, x ≠ y)

▸ 출력

총 Q개의 줄을 출력한다. 각 줄에는 입력으로 주어진 두 노드 사이의 거리를 출력한다.

 

문제 풀이  

범위랑 예외처리로 시간이 좀 걸린 문제이다. 해당 문제는 쿼리로 주어진 두 노드의 최소공통조상을 찾아서 거리를 계산해주는 문제이다. 루트노드가 주어졌기 때문에 해당 노드의 높이와 부모노드를 구하는 메서드만 구현하면 풀이가 가능하다. 단 노드의 최대 갯수가 10^15개이기 때문에 int의 범위를 초과한다. 배열의 저장공간도 10^15개를 저장할 수 없기 때문에 dp를 사용하지도 못한다. 그래서 범위 설정(long) 쿼리와 예외처리를 잘 작성해줘야 한다. 

 

노드 높이 구하기

각 노드는 k개의 자식 노드를 순서대로 갖는다. 이를 이용하면 간단하게 높이를 구할 수 있다. 주어진 노드 번호가 idx라면, k^h보다 작거나 같은 구간을 찾아주면 된다.

 

높이 구하기

 

부모노드 구하기

모든 트리의 루트노드 1로, 2부터 자식노드가 시작된다. 그리고 각 자식노드의 갯수는 k개이다. 

1번의 자식노드 3개로, (2,3,4) -2  > (0,1,2)

2번의 자식노드 3개로, (5,6,7) -2  > (3,4,5)

...

다음과 같이 2번 이후의 노드를 전부 -2씩 앞으로 당겨서  /k+1을 해주면 부모노드를 구할 수 있다.

 

부모노드 idx = (현재node idx-2)/k+1

 

 

부모 노드 구하기

 

설계

1. x, y의 각 depth를 구한다. xh, yh
   1. if(xh < yh) 로직 작성을 편하게 하기 위해 x와 y를 바꿔준다. (xh > yh)
2. if(xh!=yh) xh의 높이가 yh의 높이가 같아질 때 까지 x의 부모노드를 탐색한다. mv++
3. 높이가 같아졌으면, x와 y의 최소공통조상(LCA)를 찾는다. 두 노드가 동시에 움직이기 때문에 +2씩 카운트해준다. mv+=2

 

그리고 최악의 경우 편향트리일 경우(k==1), 위의 로직으로 실행하면 시간초과가 발생한다. 10^15의 거리를 10만번 탐색할 수도 있기 때문이다. 따라서 예외처리를 해줘야 시간초과에서 벗어날 수 있다.

 

풀이 코드 

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

public class Main {
	static long n;
	static int k, h;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		n = Long.parseLong(st.nextToken());
		k = Integer.parseInt(st.nextToken());
		int q = Integer.parseInt(st.nextToken());
		
		StringBuilder sb = new StringBuilder();
		for(int i=0; i<q; i++) {
			st = new StringTokenizer(br.readLine());
			long x = Long.parseLong(st.nextToken());
			long y = Long.parseLong(st.nextToken());
			if(k==1) {
				sb.append(Math.abs(x-y)+"\n");
			}else {
				sb.append(LCA(x,y)+"\n");
			}
		}
		System.out.println(sb.toString());
	}
	
	static long LCA(long x, long y) {
		long mv=0;
		// x,y depth
		long xh = getDepth(x);
		long yh = getDepth(y);
		
		if(xh < yh) {
			long tmp = x;
			x = y;
			y = tmp;
			long ttmp = xh;
			xh = yh;
			yh = ttmp;
		}
		
		while(xh != yh) {
			x = getParent(x);
			xh = getDepth(x);
			mv++;
		}
		
		while(x!=y) {
			x = getParent(x);
			y = getParent(y);
			mv+=2;
		}
		return mv;
	}
	
	static long getParent(long idx) {
		return (idx-2)/k+1;
	}
	
	static long getDepth(long idx) {
		if(idx==1) return 0;
		
		long line =1;
		long h =1;
		while(true) {
			line += (long)Math.pow(k,h++);
			if(idx <= line) break;
		}
		return h-1;
	}
}