boj, 백준) 11812. K진 트리 (C/C++)

1. 문제 링크

www.acmicpc.net/problem/11812

11812번: K진 트리

2. 문제 개요

각 노드가 자식을 최대 K개 가질 수 있는 트리를 K진 트리라고 한다. 총 N개의 노드로 이루어져 있는 K진 트리가 주어진다. 이전 깊이를 모두 채운 경우에만, 새로운 깊이를 만드는 것이고, 이 새로운 깊이의 노드는 가장 왼쪽부터 차례대로 추가한다.

 

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

 

 

3. 문제 힌트

이 문제는 Query가 최대 10만개이다. N도 10의 15승.. 숫자가 매우 크다. 그래서 최단경로 알고리즘인 dijkstra를 쓴다 해도.. 시간 복잡도가 너무 커진다.

 

주어진 데이터가 Tree라는 것에 주의를 기울여보자.

그리고, LCA를 통해서 거리를 풀게 될텐데, (x-LCA-y) 어떻게 LCA를 찾을 수 있을지..?

 

LCA를 찾기 위해서는 부모가 누구인지, 깊이가 얼마인지.. 등을 알아야 한다. 

기본 LCA알고리즘에서 부모를 찾을 때 시간 복잡도, 깊이를 구할 때의 시간 복잡도가 얼마였는지 생각해보자.

 

n이 매우 크기 때문에, 직접 탐색하기는 무리이고, 수학적으로 그 깊이나 부모를 계산해낼 수 없을까? 도 생각해보자.

 

다행히 이 문제는 Tree이고, 완전k진트리(?)이기 때문에 제 시간 내에 풀 수 있다.

 

4. 문제 풀이

우선, 문제를 풀기 전에 기본 LCA 알고리즘을 생각해보자.

 

처음에 root로부터의 깊이를 측정한다.

이때, DFS를 사용하는데 결국 dfs탐색은 모든 노드를 한 번씩 방문하기 때문이 O(n)이다. 그런데 이 문제에서는 n이 최대 10의 15승이기 때문에.. DFS를 사용하면 안 된다는 것을 생각할 수 있다. 또, 배열로 트리를 모두 갖고 있어도 안 된다. (공간부족)

 

우선, 계산의 편의를 위해 노드를 1부터 시작하는 것이 아닌 0부터 시작하는 것으로 변경하는 게 좋다.

 

모든 예는 K=3일 때로..

 

N=15, K=3일때의 Tree

문제의 Tree 생성 규칙에 따라 다음과 같이 나타낼 수 있다. 

 

 

자, 이제 가장 먼저 깊이를 계산해볼 것이다.

깊이는 간단하게 한 칸씩 내려가볼 건데, 내려가기전에 한칸씩 내려가도 괜찮은지 생각해보자.

 

만약 예를 들어, N=10^15(최댓값), K=1일 때를 생각해보자.

만약에 한 칸씩 내려가서 10^15의 깊이가 얼마인지 측정하려고 하면 10^15번의 반복이 필요하기 때문에 불가능이다. 그래서 K=1일 때는 예외로 처리해 주어야 한다.

 

이제, K≥2 일 때를 생각해보자,

완전K진트리이기 때문에, K=2일 때의 최대 깊이는 log_2 (10^15)가 될 것이다. 그래서 한 칸씩 내려가게 되더라도 시간 초과가 나지 않는다. 최대 log_2 (10^15)번만 탐색하면 되기 때문!

 

그럼, 해당 깊이의 가장 왼쪽 값과, 가장 오른쪽 값을 구해서, 내가 깊이를 구하려고 하는 숫자가 이 범위 내부에 속하는지 체크하면 된다.

 

깊이와 값의 범위

 

0은 예외니까 놔두고, 깊이가 1일 때, 좌 우의 범위는 위의 그림과 같다.

그다음, 깊이가 1칸씩 늘어날 때마다 현재의 값에 K를 곱하고 왼쪽은 1, 오른쪽은 K를 더하는 규칙을 찾을 수 있다. 위의 방법으로 깊이를 계산한다.

 

그리고, 이제 자신의 바로 위 부모를 찾는다. LCA알고리즘 같은 경우는 1, 2, 4, 8,... 번째 부모를 찾지만, 이 경우는 Tree의 깊이가 별로 깊지 않아(log_K (n)) 자신의 바로 위 부모만 알고 있어도 괜찮다.

 

4, 5, 6의 부모는 1이다. 13, 14, 15의 부모는 4이다. 이 규칙을 생각해보면

(x-1)/k의 내림은 부모이다.라고 구할 수 있다. 

 

 

그 뒤의 알고리즘은 LCA와 같다.

 

두 정점을 입력받아, 깊이를 같게 해준 다음에 서로가 같아질 때까지 한 칸씩 부모로 이동해나가는 것이다.

노드 간 거리는 1이기 때문에 두 노드가 동시에 올라간다면 2씩 값을 누적시켜주면 된다.

 

아 그리고 n이 커서 long long으로 해주었다.

5. 코드

#include <cstdio>
#include <algorithm>
using namespace std;

int k, q;
long long int n;

long long int get_depth(long long int val) {
	
	long long int depth = 0;

	if (k == 1) {
		depth = val;
	}
	else
	{
		if (val != 0)
		{
			depth = 1;
			long long int left = 1, right = k;

			while (!(left <= val && val <= right)) {
				++depth;
				left = left * k + 1;
				right = right * k + k;
			}
		}
	}
	return depth;
}

long long int get_parent(long long int u)
{
	long long int parent;

	parent = (u - 1) / k;

	return parent;
}

long long int get_distance(long long int u, long long int v) 
{
	long long int depth_u = 0, depth_v = 0;
	long long int ans = 0;

	depth_u = get_depth(u);
	depth_v = get_depth(v);

	if (depth_u < depth_v) {
		swap(depth_u, depth_v);
		swap(u, v);
	}
	//u is deeper than v

	if (k == 1) {
		ans = depth_u - depth_v;
	}
	else
	{
		long long int diff = depth_u - depth_v;

		ans += diff;
		for (int i = 0; i < diff; ++i) 
			u = get_parent(u);
		
		while (u != v) {
			u = get_parent(u);
			v = get_parent(v);
			ans += 2;
		}
	}

	return ans;
}

int main()
{
	scanf("%lld %d %d", &n, &k, &q);


	for (int i = 0; i < q; ++i) {
		long long int u, v;
		scanf("%lld %lld", &u, &v);

		printf("%lld\n", get_distance(u - 1, v - 1));
	}


	return 0;
}

6. 결과