[백준] 11438. LCA 2

문제 링크

[백준] 11438. LCA 2

풀이 과정

[백준] 11437. LCA와 문제는 동일하지만, N과 M의 개수가 훨씬 더 많고, 더 적은 시간 제한이 걸려 있는 문제입니다.

이 문제는 N과 M의 최대 값이 둘 다 10만이므로, 최악의 경우 100초가 넘는 시간이 소요될 것이므로 LCA 문제의 풀이를 그대로 사용하게 되면 시간 초과가 발생합니다.

최소 공통 조상을 찾기 위해, 트리를 위로 한 칸씩 탐색하는 방법이 아닌, 찾아 올라가는 횟수를 줄이기 위해선 다음과 같은 방법이 필요합니다. 예를 들어, 한 노드의 7번 째 부모를 찾기 위해서는 다음의 방법을 생각해 볼 수 있습니다.

• 파란색 경로 : 한 칸씩 부모를 조상을 거슬러 올라감.
• 빨간색 경로 : 4칸을 먼저 올라간 뒤, 2칸, 1칸을 순서대로 올라감.

위와 같은 방법을 이용해 부모의 부모를 거듭해 올라가는 작업을 단축시킬 수 있습니다.

단순히 각 노드의 부모만을 저장하는 방식이 아닌, 노드의 조상들의 번호를 저장합니다. 시간상, 메모리상 모든 조상들을 다 저장할 필요는 없으므로, 2의 거듭제곱 번째 조상들만 저장합니다.

정점의 개수가 최대 10만개 이므로, 편향 트리 구조 같이 최악의 경우, 가장 깊이가 깊은 리프 노드의 조상 개수는 (10만 - 1)개가 됩니다. $2^{16}=65,536$ , $2^{17}=131,072$ 이므로, 18개의 인덱스로 모든 조상들을 표현 가능하므로, 조상을 저장하는 ancestor 변수는 (N+1) x 18 의 크기를 갖습니다.

ancestor = new int[N + 1][18];

dfs 탐색을 하며, 트리를 구성하는 모든 노드들의 깊이와 부모를 저장합니다. 부모는 $2^{0}$ 번째 조상을 의미하며, 따라서 0번 인덱스에 부모 노드 번호를 입력합니다.

ancestor[now][0] = parent;

다음으로, 각 노드별 모든 조상들의 번호를 구합니다. 로직은 다음과 같습니다.

• 16번 노드의 $2^1$번째 조상 = 16번 노드의 $2^0$번째 조상의 $2^0$번째 조상
• 16번 노드의 $2^2$번째 조상 = 16번 노드의 $2^1$번째 조상의 $2^1$번째 조상
• 16번 노드의 $2^3$번째 조상 = 16번 노드의 $2^2$번째 조상의 $2 ^2$번째 조상

위 로직을 코드로 적으면 아래와 같습니다.

for (int j = 1; j <= 17; j++) {
    for (int i = 1; i <= N; i++) {
        ancestor[i][j] = ancestor[ancestor[i][j - 1]][j - 1];
    }
}

다음으로, lca() 메소드는 아래와 같이 세 동작으로 이뤄집니다.

1. 두 노드 a, b를 받아 깊이가 더 깊은 노드를 b로 설정합니다.
2. 깊이가 다르면, 두 노드의 깊이를 맞춥니다.
3. 같은 깊이의 두 노드를 가지고 최소 공통 조상을 찾아냅니다.

아래는 두 노드의 깊이를 일치시키는 코드입니다. 두 노드의 깊이 차를 구한 다음, 다음 사진과 같이 동작하게 됩니다. 예를 들어, 두 노드의 깊이 차가 5라면, 5는 1+4 로 표현되니, 순서대로 한칸 → 네칸을 올리면 깊이가 맞춰집니다.

if (depth[a] != depth[b]) {
    int diff = depth[b] - depth[a];

    for (int i = 0; i <= 17; i++) {
        if ((diff & (1 << i)) != 0) {
            b = ancestor[b][i];
        }
    }
}

깊이가 맞춰졌다면, 2의 거듭제곱으로 표현되는 올릴수 있는 가장 큰 수만큼 올려가면서 최소 공통 조상을 찾습니다. 예를 들어, 아래 사진과 같이 노드 a, b의 9번째 조상이 최소 공통 조상인 경우, $2^3=8$ 칸을 올려 a와 b로 위치시킵니다. 그 위로는 노드 a와 b의 조상들 번호들이 다 같으므로 작업을 멈춥니다. 멈춘 위치 노드 a와 b의 첫 번째 조상(=부모)인 ancestor[a][0] 와 ancestor[b][0] 값은 같으므로, 둘 중 하나를 반환하면 문제를 해결할 수 있습니다.

for (int i = 17; i >= 0; i--) {
    if (ancestor[a][i] != ancestor[b][i]) {
        a = ancestor[a][i];
        b = ancestor[b][i];
    }
}

return ancestor[a][0];

코드

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Scanner;

public class Main {
    static int N, M;
    static ArrayList<Integer>[] adj;
    static boolean[] visit;
    static int[] depth;
    static int[][] ancestor;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        StringBuilder sb = new StringBuilder();
        N = sc.nextInt();

        adj = new ArrayList[N + 1];
        for (int i = 1; i <= N; i++)
            adj[i] = new ArrayList<>();
        visit = new boolean[N + 1];
        depth = new int[N + 1];
        ancestor = new int[N + 1][18];

        for (int i = 0; i < N - 1; i++) {
            int v1 = sc.nextInt();
            int v2 = sc.nextInt();

            adj[v1].add(v2);
            adj[v2].add(v1);
        }

        dfs(1, 1, 0);
        setAncestor();

        M = sc.nextInt();
        for (int i = 0; i < M; i++) {
            sb.append(lca(sc.nextInt(), sc.nextInt()) + "\n");
        }

        System.out.println(sb);
    }

    private static void dfs(int parent, int now, int d) {
        visit[now] = true;
        depth[now] = d;
        ancestor[now][0] = parent;

        for (int next : adj[now]) {
            if (!visit[next]) {
                dfs(now, next, d + 1);
            }
        }
    }

    private static void setAncestor() {
        for (int j = 1; j <= 17; j++) {
            for (int i = 1; i <= N; i++) {
                ancestor[i][j] = ancestor[ancestor[i][j - 1]][j - 1];
            }
        }
    }

    private static int lca(int a, int b) {
        if (depth[a] > depth[b]) {
            int tmp = a;
            a = b;
            b = tmp;
        }

        if (depth[a] != depth[b]) {
            int diff = depth[b] - depth[a];

            for (int i = 0; i <= 17; i++) {
                if ((diff & (1 << i)) != 0) {
                    b = ancestor[b][i];
                }
            }
        }

        if (a == b) return a;

        for (int i = 17; i >= 0; i--) {
            if (ancestor[a][i] != ancestor[b][i]) {
                a = ancestor[a][i];
                b = ancestor[b][i];
            }
        }

        return ancestor[a][0];
    }
}