컴퓨터 사이언스/1고리즘

백준 11438 : LCA 2

저세상 개발자 2021. 8. 21. 02:27

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

 

11438번: LCA 2

첫째 줄에 노드의 개수 N이 주어지고, 다음 N-1개 줄에는 트리 상에서 연결된 두 정점이 주어진다. 그 다음 줄에는 가장 가까운 공통 조상을 알고싶은 쌍의 개수 M이 주어지고, 다음 M개 줄에는 정

www.acmicpc.net

 

LCA 1과 문제 자체는 같으나 N의 크기가 10만, 쿼리 M의 크기가 10만으로 더 크다. 따라서 O(N * M)은 무조건 시간초과가 나므로 O(MlogN)으로 해결해야하는 문제다.

LCA 알고리즘도 DP를 이용해서 메모이제이션 하면서 조건을 만족하는 부모노드를 이분탐색 비슷한 개념으로 O(logN)의 시간복잡도로 찾을 수 있다.

 

아래 링크를 보고 공부했으며, 아래 링크에 너무너무 잘 설명되어있다.

https://www.crocus.co.kr/660

 

Crocus

Beginner와 Developer사이의 Crocus

www.crocus.co.kr

 

오늘도 내 나름대로 공부한 것을 정리해보자면 아래와 같다.

 

LCA 알고리즘은 O(n) 풀이도, O(logN)풀이도 공통적으로 두 단계로 나뉜다.

 

1. 최소 공통 조상을 찾고자 하는 두 노드를 A, B라고 한다면, A와 B의 깊이부터 맞춘다.

   -> 깊이가 더 큰 노드를 깊이가 더 작은 노드의 깊이와 같도록 이동시킨다.

2. A와 B의 깊이가 같은 상태이므로, A와 B를 같은 변위만큼 이동시키면서 최소 공통 조상을 찾는다.

 

O(N) 풀이는 1, 2번 작업을 수행할 때 깊이를 1칸씩 선형으로 이동한다.

O(logN) 풀이는 2^i 칸 씩 이동한다. 예를 들어 3칸을 위로 가야 한다면 2^1 칸을 움직인 후 움직인 위치를 기준으로 다시 2^0 칸을 움직이는 원리다.

만약 이게 잘 와닿지 않는다면 다른 예시를 들어보자.

예를 들어 1,024칸을 이동한다면 logN풀이는 1번에 이동하고 N풀이는 1,024번만에 이동한다고 생각하면 된다.

(사실 이동 가능한지 확인하는 연산이 추가되기 때문에 정확히 1번이라고 하기는 그렇지만 N풀이보다는 훨씬 더 적은 연산이 필요하다는 뜻으로 받아들여주면 좋겠다.)

 

이를 위해 각 노드들의 2^i 번째 부모 노드를 미리 메모이제이션 해둔다.

처음에 bfs든, dfs든 사용해서 2^0번 째(첫 번째 부모) 노드 값을 저장했다면, 아래와 같은 로직으로 각 노드들의 2^i 번째 부모 노드를 저장할 수 있다.

for (int j = 1; j < MAX_POW; j++) {
	for (int i = 1; i <= n; i++) {
		dp[i][j] = dp[dp[i][j - 1]][j - 1];
	}
}

여기서 MAX_POW는 최악의 경우(트리가 1렬로 늘어서는 경우)에 리프노드가 부모노드 정보를 모두 저장하기 위해서 필요한 값이다. 즉, 이 문제에서는 노드의 개수가 최대 10만개이고, 2^17 = 131,072로 최초로 10만이 넘어가는 수 이므로, MAX_POW는 18로 줬다. (인덱스 17을 사용하기 위해서)

 

A와 B의 depth를 맞추는 1번 로직은 아래와 같이 구현할 수 있다. B가 현재 깊이가 더 큰 상황이고 A와 같은 깊으로 맞추기 위하여 부모 노드로 이동시키는 과정이다.

// a와 b의 depth 맞추기
if (depth[a] != depth[b]) {
	// 가장 먼 부모노드부터
	for (int i = MAX_POW - 1; i >= 0; i--) {
		// 만약 현재 노드의 2^i번 째 부모노드의 깊이가 a의 깊이보다 작다면 continue
		// 현재 노드의 2^i번 째 부모가 존재하지 않을 경우 노드 번호 0번
		if (depth[dp[b][i]] < depth[a]) continue;
		// 현재 노드의 2^i번 째 부모노드의 깊이가 a의 깊이 보다 큰 경우
		// 현재 노드를 현재 노드의 2^i번 째 부모노드로 바꿔줌
		// 그리고 계속 for문을 돌면서 바뀐 현재 노드를 기준으로 
		// 2^(i - 1)번째 부모노드부터 동일 로직 수행
		// 이렇게 함으로써 마치 이분탐색처럼 범위를 줄여가면서
		// O(logN)의 시간복잡도로 a의 깊이와 같은 노드로 이동 가능
		b = dp[b][i];
	}
}

주석에도 적혀있지만 다시 적어보자면,

i는 MAX_POW부터 시작해서, 즉, 현재 노드의 가장 먼 부모부터 시작해서 해당 노드의 깊이와 A의 깊이를 비교한다.

A의 깊이보다 더 작다면 A와 B의 깊이를 맞춘다는 목적에 어긋나므로 continue해서 i를 증가시켜 깊이가 더 깊은 B의 부모를 검색한다.

만약 현재 i값에서 B의 2^i번 째 부모 노드의 깊이가 A의 깊이보다 작다면 B를 현재 B의 2^i번 째 부모 노드로 설정하고 i를 1 감소시켜 다시 동일 로직을 수행한다.

dp[b][i] = dp[dp[b][i-1]][i-1] 이므로, i가 현재 i - 1보다 큰 값에 대해서는 검사할 필요가 없다.

(이미 B를 바꾸기 전에 B에서 검사했고, A의 깊이보다 작아서 continue 되었다.)

 

2번 로직도 1번 로직을 구현한 것과 유사하게 구현할 수 있다.

// 다시 최대 제곱수부터
for (int i = MAX_POW - 1; i >= 0; i--) {
	// a와 b의 현재 깊이는 같고
	// 특정 깊이(최소 공통 조상의 높이) 이하의 깊이에서는 항상 같은 조상을 가질 것임
	// 따라서 조상이 같은 동안은 continue
	if (dp[a][i] == dp[b][i]) continue;
	// 조상이 다르다면 그 서로 다른 조상으로 a와 b를 옮겨줌
	// 옮겨준 후 계속 로직 수행, depth 맞출때와 마찬가지임
	a = dp[a][i];
	b = dp[b][i];
}
// 결과적으로 a와 b는 최소 공통 조상의 첫 번째 자식들로 바뀜

 

전체 코드는 아래와 같다.

#include <iostream>
#include <vector>
#include <queue>

using namespace std;

const int MAX = 100'001;
const int MAX_POW = 18; // 2^17=131,072

int n;
vector<int> edge[MAX];
int depth[MAX];
int dp[MAX][MAX_POW];

// 트리 만들기 -> depth, dp[i][0] (각 노드 i의 첫 번째 부모 노드) 저장
void make_tree() {
	int cur, next;
	int cur_height = 0, cur_height_size = 0;
	queue<int> q;
	q.push(1);
	dp[1][0] = 1;

	// 0번 노드는 실제로 존재하는 노드는 아니지만
	// 어떤 노드의 2^i번 째 부모 노드가 존재하지 않는다면 그 값을 0으로 세팅해 줄 것임
	// 그리고 이후 LCA알고리즘에서 부모 노드의 depth를 비교하는 로직이 있으므로
	// depth[0]은 루트 노드보다도 낮은 depth인 -1로 설정
	depth[0] = -1;

	while (!q.empty()) {
		if (cur_height_size <= 0) {
			cur_height_size = q.size();
			++cur_height;
		}
		cur = q.front();
		q.pop();

		for (int i = 0; i < edge[cur].size(); i++) {
			next = edge[cur][i];
			if (dp[next][0]) continue;

			dp[next][0] = cur;
			depth[next] = cur_height;
			q.push(next);
		}

		--cur_height_size;
	}
	dp[1][0] = 0;
}

// dp[i][j] -> 각 노드 i의 2^j 번 째 부모노드 저장, 미리 저장한 부모 노드의 부모노드 정보(dp값)를 사용 
void update_dp() {
	for (int j = 1; j < MAX_POW; j++) {
		for (int i = 1; i <= n; i++) {
			// 현재 노드의 2^j번 째 부모는 2^(j-1)번 째 부모노드의 2^(j-1)번 째 부모 노드
			dp[i][j] = dp[dp[i][j - 1]][j - 1];
		}
	}
}

int LCA(int a, int b) {
	// b의 depth값을 더 크게 만들기 위하여 swap
	if (depth[a] > depth[b]) { 
		int tmp = a;
		a = b, b = tmp;
	}

	// a와 b의 depth 맞추기
	if (depth[a] != depth[b]) {
		// 가장 먼 부모노드부터
		for (int i = MAX_POW - 1; i >= 0; i--) {
			// 만약 현재 노드의 2^i번 째 부모노드의 깊이가 a의 깊이보다 작다면 continue
			// 현재 노드의 2^i번 째 부모가 존재하지 않을 경우 노드 번호 0번
			if (depth[dp[b][i]] < depth[a]) continue;
			// 현재 노드의 2^i번 째 부모노드의 깊이가 a의 깊이 보다 큰 경우
			// 현재 노드를 현재 노드의 2^i번 째 부모노드로 바꿔줌
			// 그리고 계속 for문을 돌면서 바뀐 현재 노드를 기준으로 
			// 2^(i - 1)번째 부모노드부터 동일 로직 수행
			// 이렇게 함으로써 마치 이분탐색처럼 범위를 줄여가면서
			// O(logN)의 시간복잡도로 a의 깊이와 같은 노드로 이동 가능
			b = dp[b][i];
		}
	}

	// depth 맞추기 예시
	// 현재 depth[a] = 1, depth[b] = 4인 상황이라면
	// depth를 맞추기 위해 b가 자신보다 3단계 위에 있는 부모 노드가가 되어야함
	// 위의 for문을 돌다가 
	// i = 2일 때, depth[dp[b][i]] = 0 으로 depth[a]보다 작으므로 continue
	// i = 1일 때, depth[dp[b][i]] = 2 로 depth[a]보다 작지 않으므로 b = dp[b][i] -> depth[b] = 2가 됨
	// 바뀐 b에 대하여 다시 부모노드를 확인해야하는데 i는 다시 -1된 i = 0 부터 검사
	// dp[b][i] = dp[dp[b][i-1]][i-1]이므로 
	// 현재 상황에서 i가 1 이상인 경우는 이미 b가 바뀌기 전에 검사한 i가 2인 경우와 같은 경우
	// 따라서 i가 1 줄어든 i = 0부터 검사하면 됨
	// i = 0일 때, depth[dp[b][i]] = 1 로 depth[a]보다 작지 않으므로 b = dp[b][i] -> depth[b] = 1가 됨
	// 따라서 depth[a] = depth[b]가 됨


	// 이후 로직에서 a와 b를 최소 공통 조상의 자식노드로 바꿔줄 것이므로
	// a == b인 경우는 아무것도 바뀌지 않고 그대로 로직을 빠져나와 
	// a의 첫 번째 부모를 리턴해 줄 것임, 따라서 미리 체크
	if (a == b) return a;

	// 다시 최대 제곱수부터
	for (int i = MAX_POW - 1; i >= 0; i--) {
		// a와 b의 현재 깊이는 같고
		// 특정 깊이(최소 공통 조상의 높이) 이하의 깊이에서는 항상 같은 조상을 가질 것임
		// 따라서 조상이 같은 동안은 continue
		if (dp[a][i] == dp[b][i]) continue;
		// 조상이 다르다면 그 서로 다른 조상으로 a와 b를 옮겨줌
		// 옮겨준 후 계속 로직 수행, depth 맞출때와 마찬가지임
		a = dp[a][i];
		b = dp[b][i];
	}
	// 결과적으로 a와 b는 최소 공통 조상의 첫 번째 자식들로 바뀜

	// 최소 공통 조상 리턴
	return dp[a][0];
}

void solution() {
	int m;
	int a, b;
	cin >> n;
	for (int i = 1; i < n; i++) {
		cin >> a >> b;
		edge[a].push_back(b);
		edge[b].push_back(a);
	}

	make_tree();
	update_dp();

	cin >> m;
	while (m--) {
		cin >> a >> b;
		cout << LCA(a, b) << '\n';
	}
}

int main() {
	ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
	solution();

	return 0;
}
메모리: 14960 kb 시간: 148 ms

 

여담)

logN 풀이는 정말 심오한 것 같다. 아이디어도 좋지만, 아이디어를 구현하는 방법도 정말 군더더기 없이 깔끔하다.

logN 코드를 참고하기 전에 logN이 되는 이유와 로직만 글로 먼저 이해하고 직접 구현해봤는데 내 코드는 이중포문으로 끔찍히 가독성이 좋지 않은 코드였다. 정답은 for문 하나로 처리하는 것 보고 감탄이 절로 나왔다. 알고리즘 뿐만 아니라 아직 구현력도 많이 부족함을 느끼게 해준 문제였다.