[백준/C++] 10838. 트리

문제

https://www.acmicpc.net/problem/10838 

N개의 노드로 구성된 루트가 있는 트리가 다음과 같이 주어진다. 각 노드는 0부터 N-1까지의 번호로 구별되고, 0번 노드는 루트 노드이고, 나머지 노드 각각은 0번 노드의 자식 노드이다.

트리에 적용할 수 있는 연산은 세 종류이며, 이를 통해 트리의 모양을 바꾸거나 트리 에지에 색칠을 할 수 있다. 각 연산과 그 의미는 다음과 같다. 

1. paint(a, b, c): a번 노드와 b번 노드를 잇는 최단 경로를 찾은 후, 경로 상에 있는 모든 에지를 색깔 c로 칠한다. 
2. move(a, b): a번 노드의 부모 노드를 b번 노드로 바꾼다. 단, b번 노드는 a번 노드를 루트로 하는 부트리(subtree)에 속하지 않는다. 부모 노드를 바꾸기 전 a번 노드의 부모 노드를 p라 할 때, 새로운 에지 (a,b)는 원래의 에지 (a,p)의 색깔을 갖는다. 
3. count(a, b): a번 노드와 b번 노드를 잇는 최단경로를 찾은 후, 그 경로 사이에 있는 에지에 칠해진 서로 다른 색깔의 개수를 출력한다. 

에지에 칠하는 색깔 c를 정수로 표시한다. 그리고 처음에는 모든 에지의 색깔이 0이라고 가정한다. 

풀이

paint명령과 count명령에서 노드간 최단경로를 찾아야 한다. 

트리에서는 두 노드간 최단경로는 무조건 두 노드의 가장 가까운 공통 조상을 거쳐간다! 

따라서 지난번에 풀었던 가장 가까운 공통 조상 풀이를 활용하면 된다고 생각했다. 

 

parent배열에 pair<int,int>(부모 노드 번호, 색깔) 을 저장해서 이용하였다. 

 

가장 가까운 공통 조상을 찾는 알고리즘을, 지난번에 풀었던 방식으로 했더니 시간초과가 났다! 

 

검색해보니 LCS(Least Common Ancestor) 이라는 이름으로 꽤나 정형화된? 유명한? 알고리즘이었다. 

세그먼트 트리로도 풀 수 있다고 하는데, 지금은 이해가 안되어서 DP 비슷한 방식을 이용했다.

 

check배열을 이용해서, a와 b사이의 가장 가까운 공통 조상을 찾을 때, 한 단계씩 거슬러 올라가면서 발견한 노드를 체크배열에 표시한다. 표시하기전에, 이미 표시된 노드이면 그 노드가 가장 가까운 공통 조상이 된다!  

 

check배열을 bool로 하고, 매번 공통조상을 찾을 때 초기화하게 했더니 시간초과가 나는 것 같아서, 

int로 하고 매번 공통조상을 찾을 때 마다 표시하는 값을 하나씩 증가시켜서 사용했다. 

코드

#include <iostream>
#include <vector>
#include <unordered_set>

#define sc scanf_s
using namespace std;
int N, K;
pair<int, int> parent[100001]; // <부모 노드 번호, 색깔>
int check[100001];

int temp = 0;
int findCommonParent(int a, int b) {
	if (a == b) return a;
	int cnt = 0;
	temp++;
	while (cnt <= 1000) {
		if (a != 0) {
			if (check[a] == temp)
				return a;
			check[a] = temp;
			a = parent[a].first;
		}
		if (b != 0) {
			if (check[b] == temp)
				return b;
			check[b] = temp;
			b = parent[b].first;
		}
		if (a == b && a == 0) return 0;
		cnt++;
	}
}
int count(int a, int b) { 
	unordered_set<int> colors;
	int commonP = findCommonParent(a, b);
	while (a != commonP) {
		colors.insert(parent[a].second);
		a = parent[a].first;
	}
	while (b != commonP) {
		colors.insert(parent[b].second);
		b = parent[b].first;
	}
	return (int)colors.size(); 
}
void paint(int a, int b, int c) { 
	int commonP = findCommonParent(a, b);
	while (a != commonP) {
		parent[a].second = c;
		a = parent[a].first;
	}
	while (b != commonP) {
		parent[b].second = c;
		b = parent[b].first;
	}
}
void func(int op) {
	int a, b; sc("%d %d", &a, &b);
	if (op == 1) { //paint 
		int c; 
		sc("%d", &c);
		paint(a, b, c);
	}
	else if (op == 2) { // move
		parent[a].first = b; //a의 부모를 b로 바꾼다! 
	}
	else { // count
		int cnts = count(a, b);
		printf("%d\n", cnts);
	}
}
int main() {
	sc("%d %d", &N, &K);
	parent[0].first = -1;
	for (int k = 0;k < K;k++) {
		// 연산의 종류 : (1, paint) (2, move) (3, count) 
		int op; 
		sc("%d", &op);
		func(op);
	}
	return 0;
}

결과