SWEA 1855 <-클릭


LCA알고리즘을 사용하여 최소공통 조상을 구하는 알고리즘이다. 처음에 LCA알고리즘을 몰라서 한칸씩 부모를 타고 올라가 시간 초과가 났던 문제이다.

LCA란?

아이디어를 간단하게 요약하면 다음과 같다

• 트리 구조체에서 자신의 2의 거듭제곱번째 조상을 dp배열에 저장한다.
  
  
• 특정한 두 노드의 최소 공통 조상을 찾을 때 dp배열을 활용하여 1칸->2칸->4칸 ...씩 타고 올라간다.

일단 dp배열을 생성하는 코드는 다음과 같다

for (int j = 1; j < floor(log2(depth[i])) + 1; j++) {//1번 부터는 dp
	dp[i][j] = dp[dp[i][j - 1]][j - 1];
}

i 의 $2^j$번째 조상은 i의 $2^j-1$번째 조상의 $2^j-1$번째 조상이다 라고 이해하면 된다.

영준이가 방문할 노드는 BFS를 따라가기 때문에 Queue에 방문 순서를 넣어 이를 구현하였다

Q.push(1);
int front;
int cur;
while (!child[front].empty()) { //자식들 푸쉬
	Q.push(child[front].front());
	child[front].pop();
}

이제 LCA를 사용하여 두 노드의 최소 공통 조상을 찾아야 한다.

이 문제에서 두 노드는 같은 층에 있거나 1층 차이가 나기 때문에 두 케이스를 분류하였다.

if (depth[cur] != depth[front]) { //층이 다른경우
	cur = parent[cur]; //층을 맞춰줌
	flag = true;
}

flag가 1인 경우에는 층을 맞추기 위해 cur노드를 한칸 위로 올린 경우이므로 나중에 +1을 해준다.

while (cur != front) {
	if (depth[cur] < pow(2, iter)) {
		iter = 0;
	} //root를 넘어가는 경우 iter = 0
	if (dp[cur][iter] != dp[front][iter]) {//이번 iter에도 다른 경우
		cur = dp[cur][iter];
		front = dp[front][iter];
	}
	else if (iter == 0) { //이번 iter에 같아졌는데 iter가 0인 경우
		cur = dp[cur][0];
		front = dp[front][0];
	}
	else { //이번 iter에 같아지고 iter가 0이 아닌 경우
		iter = -1;
	}
iter++;
}
if (flag) cnt += (2 * (depth[save_cur] - depth[cur])) + 1; //층이 달랐던 경우
else cnt += (2 * (depth[save_cur] - depth[cur])); // 층이 같았던 경우

LCA알고리즘을 사용하여 최소 공통 조상을 구하는 부분의 코드는 위와 같다.

코드

github

#define _CRT_SECURE_NO_WARNINGS

#include<iostream>
#include<queue>
#include<cmath>

using namespace std;

int parent[100002];
int depth[100002];
int dp[100002][16];
int N;
long long cnt = 0;
queue<int> child[100002];
queue<int> Q;

void solve();

int main(int argc, char** argv)
{
	int test_case;
	int T;
	freopen("input/1855_input.txt", "r", stdin);
	cin >> T;
	for (test_case = 1; test_case <= T; ++test_case)
	{
		cout << "#" << test_case << " ";
		cin >> N;
		int input;
		depth[1] = 0;
		for (int i = 2; i <= N; i++) {
			cin >> input;
			parent[i] = input;
			child[input].push(i); //자식 저장
			depth[i] = depth[input] + 1;
			dp[i][0] = input;//0번에는 부모 저장
			for (int j = 1; j < floor(log2(depth[i])) + 1; j++) {//1번 부터는 dp
				dp[i][j] = dp[dp[i][j - 1]][j - 1];
			}

		}
		solve();
		cout << cnt << endl;
		cnt = 0;
	}
	return 0;//정상종료시 반드시 0을 리턴해야합니다.
}

void solve() {
	Q.push(1);
	int front;
	int cur;
	while (!Q.empty()) {
		front = Q.front();
		while (!child[front].empty()) { //자식들 푸쉬
			Q.push(child[front].front());
			child[front].pop();
		}
		Q.pop();
		if (Q.empty()) break;
		cur = Q.front();
		bool flag = false;
		if (depth[cur] != depth[front]) { //층이 다른경우
			cur = parent[cur]; //층을 맞춰줌
			flag = true;
		}
		int save_cur = cur; //저장
		int iter = 0;
		while (cur != front) {
			//printf("cur: %d, front: %d, iter: %d\n", cur, front, iter);
			if (depth[cur] < pow(2, iter)) {
				iter = 0;
			} //root를 넘어가는 경우 iter = 0
			if (dp[cur][iter] != dp[front][iter]) {//이번 iter에도 다른 경우
				cur = dp[cur][iter];
				front = dp[front][iter];
			}
			else if (iter == 0) { //이번 iter에 같아졌는데 iter가 0인 경우
				cur = dp[cur][0];
				front = dp[front][0];
			}
			else { //이번 iter에 같아지고 iter가 0이 아닌 경우
				iter = -1;
			}
			iter++;
		}
		if (flag) cnt += (2 * (depth[save_cur] - depth[cur])) + 1; //층이 달랐던 경우
		else cnt += (2 * (depth[save_cur] - depth[cur])); // 층이 같았던 경우

	}
}

정답~.~