백준 11438 : LCA 2
https://www.acmicpc.net/problem/11438
LCA 1과 문제 자체는 같으나 N의 크기가 10만, 쿼리 M의 크기가 10만으로 더 크다. 따라서 O(N * M)은 무조건 시간초과가 나므로 O(MlogN)으로 해결해야하는 문제다.
LCA 알고리즘도 DP를 이용해서 메모이제이션 하면서 조건을 만족하는 부모노드를 이분탐색 비슷한 개념으로 O(logN)의 시간복잡도로 찾을 수 있다.
아래 링크를 보고 공부했으며, 아래 링크에 너무너무 잘 설명되어있다.
오늘도 내 나름대로 공부한 것을 정리해보자면 아래와 같다.
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문 하나로 처리하는 것 보고 감탄이 절로 나왔다. 알고리즘 뿐만 아니라 아직 구현력도 많이 부족함을 느끼게 해준 문제였다.