[백준 11812] K진 트리

문제 링크 : 11812번: K진 트리 (acmicpc.net)

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개의 줄을 출력한다. 각 줄에는 입력으로 주어진 두 노드 사이의 거리를 출력한다.

---

최소 공통 조상을 이용한 문제이다. 

이 문제의 핵심은 N 값이 굉장히 커서 배열을 이용한 풀이가 금지된다. (N이 최대 \(10^{15}\) 라는점)

 

※ 문제 접근 방법

① 부모와 자식이 어떠한 관계를 갖는지 규칙성을 찾아야 한다.

이때, 문제에서 "이 새로운 깊이의 노드는 가장 왼쪽부터 차례대로 추가 한다" 라는 조건을 활용한다면 부모와 자식이 어떠한 관계를 갖는지 알 수 있다.

 

 

위 트리는 K진 트리이다. (이때, N = (K + 1) * K + 2)

여기서 부모를 x, 자식을 A = {\(a_{1}\), ... , \(a_{K}\)} 라고 하자. 

예를 들어서 x = 3 이라고 한다면, A = {\(a_{1}\), ... , \(a_{K}\)} = {2K + 2, 2K + 3, ... , 3K + 1} 이다.

이때 A 의 모든 원소에 2를 빼면 다음과 같다.

 

A = {2K, 2K + 1, 2K + 2, ..., 3K - 1} = {2K, 2K + 1, 2K + 2, ... , 2K + (K - 1)}

 

그리고 A의 모든 원소에 K를 더하면 다음과 같다.

 

A = {3K, 3K + 1, 3K + 2, ... , 3K + (K - 1)}

 

마지막으로 A의 모든 원소에 K를 나누면 다음과 같다.

 

A = {\(\frac{3K}{K}\), \(\frac{3K + 1}{K}\), \(\frac{3K + 2}{K}\), ... , \(\frac{3K + (K - 1)}{K}\)} = {3, 3, 3, ..., 3}

 

따라서 부모와 자식의 관계는 x = \(\frac{a_{i} + K - 2}{K}\) 임을 유추할 수 있다. (1 ≤ i ≤ K)

 

② ① 에서 부모와 자식의 관계를 이용하여 최소 공통 조상(LCA - Least Common Ancestor)를 구한다.

노드 x, y 에 대한 최소 공통 조상을 구하는 코드는 다음과 같다.

 

ll get_lca(ll x, ll y) {
	map<ll, bool> m;
	m[x] = m[y] = true;
	while (x != 1) {
		x = (x + K - 2) / K;
		if (m[x]) return x;
		m[x] = true;
	}
	while (y != 1) {
		y = (y + K - 2) / K;
		if (m[y]) return y;
	}
}

 

③ 최소공통조상을 lca라고 할 때, (x에서 lca까지의 거리) + (y에서 lca까지의 거리) 가 결국 정답이 된다.

임의의 노드에서 lca까지의 거리를 반환하는 함수는 다음과 같다.

 

int go_lca(ll x, ll lca) {
	int dist = 0;
	while (x != lca) {
		x = (x + K - 2) / K;
		dist++;
	}
	return dist;
}

 

④ 이때 한가지 간과한 사실이 있다.

과연 K 값이 1일 경우 시간 초과 없이 진행이 가능한지에 대한 의문이다.

여기서 K = 1 인 경우와 K ≥ 2 인 경우에 대한 트리의 depth 가 어떠한지를 알아야 한다.

 

K = 1 인 경우의 트리의 depth는  다음과 같다. 

 

 

• k = 1 인 경우 : 자식이 무조건 하나 존재하므로 트리의 depth는 N이 된다. 
• 최악의 경우 O(N)이므로 시간 내에 해결할 수 없다. (N = \(10^{15}\))
• 따라서 K = 1 인 경우, |x - y| 를 통해 답을 구해야만 한다. 

 

K ≥ 2 인 경우의 트리의 depth는  다음과 같다. (이때 K = 2 라고 가정)

 

 

• K ≥ 2 인 경우 : 자식이 K개 존재하므로 트리의 depth는 \(log_{K}N\) 이 된다.
• 최악의 경우 O(\(log_{K}N\)) 이므로 시간 내에 해결할 수 있다. (N = \(10^{15}\))
• 따라서 K ≥ 2 인 경우 최소 공통 조상을 이용해여 해결해야만 한다.

 

Solution Code

 

#define _CRT_SECURE_NO_WARNINGS
#include <string>
#include <stdio.h>
#include <math.h>
#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
#include <queue>
#include <stack>

#define INF 1000000000
#define endl '\n'
#define ll long long

using namespace std;

ll N, K, Q;

void input() {
	cin >> N >> K >> Q;
}

ll get_lca(ll x, ll y) {
	map<ll, bool> m;
	m[x] = m[y] = true;
	while (x != 1) {
		x = (x + K - 2) / K;
		if (m[x]) return x;
		m[x] = true;
	}
	while (y != 1) {
		y = (y + K - 2) / K;
		if (m[y]) return y;
	}
}

int go_lca(ll x, ll lca) {
	int dist = 0;
	while (x != lca) {
		x = (x + K - 2) / K;
		dist++;
	}
	return dist;
}

void solution() {
	ll x, y, lca;
	while (Q--) {
		cin >> x >> y;
		if (K == 1) {
			cout << abs(x - y) << endl;
		}
		else {
			lca = get_lca(x, y);
			cout << go_lca(x, lca) + go_lca(y, lca) << endl;
		}
	}
}

int main()
{
	ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);

	input();
	solution();

	return 0;
}