[백준] 11438번 LCA 2 (JAVA)

문제

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

11438번: LCA 2

 

설명

트리의 각 정점은 1번부터 N번까지 번호가 매겨져 있으며, 루트는 1번이다.
두 노드의 쌍 M개가 주어졌을 때, 두 노드의 가장 가까운 공통 조상이 몇 번인지 출력해야 한다.

각 정점의 방문 여부를 체크하기 위한 check 배열과 각 정점의 깊이를 저장할 depth 배열을 만든다.

그리고 각 정점의 \(2^k\) 번째 부모를 저장하기 위해 parent[18][] 배열을 만든다.

여기서 \(2^{17} > N = 100,000\) 이기 때문에 배열의 크기를 18로 설정했다.

마지막으로 각 정점의 연결관계를 저장할 adjList도 만들어 입력값들을 저장한다.

boolean[] check = new boolean[N + 1];
int[] depth = new int[N + 1];
int[][] parent = new int[18][N + 1];

List<List<Integer>> adjList = new ArrayList<>();
for (int i = 0; i <= N; i++) {
    adjList.add(new ArrayList<>());
}

for (int i = 1; i < N; i++) {
    st = new StringTokenizer(br.readLine());
    int v1 = Integer.parseInt(st.nextToken());
    int v2 = Integer.parseInt(st.nextToken());
    
    adjList.get(v1).add(v2);
    adjList.get(v2).add(v1);
}

이제 depth 배열을 채우기 위해 각 정점의 깊이를 구해야 한다.

먼저 루트 노드인 1번 정점을 방문한다.

방문한 정점은 방문 처리를 하고 깊이를 저장한다.

해당 정점에 인접한 정점들 중 방문하지 않은 정점을 찾아 깊이를 하나 늘린 상태로 getDepth 함수를 호출하여 DFS를 진행한다.

void getDepth(int cur, int dep) {
    // 한번 방문한 곳은 방문하지 않음
    check[cur] = true;
    // 현재 정점에 depth를 저장한다
    depth[cur] = dep;

    // 인접리스트 (adjList.get(cur).size() -> cur 점의 degree)
    for (int i = 0; i < adjList.get(cur).size(); i++) {
        int dest = adjList.get(cur).get(i);

        // 방문하지 않은 점일 경우
        if (!check[dest]) {
            // dest 점의 직계조상은 cur이다
            parent[0][dest] = cur;
            // DFS 진행
            getDepth(dest, dep + 1);
        }
    }
}

이제 모든 정점의 깊이와 직계 조상을 각각 depth 배열과 parent 배열에 저장했다.

이제 최소 공통 조상을 구하는 방법인 이분 탐색을 위해 \(2^k\) 번째 부모까지 저장한다.

// 2^17까지 조상을 구해놓는다면 배열 값이 부족할 일이 없음 (2^17 > N = 100,000)
for (int i = 1; i <= 17; i++) {
    for (int j = 1; j <= N; j++) {
        // j의 2^i번째 부모는 j의 2^(i-1)번째 부모의 2^(i-1)번째 부모이다
        parent[i][j] = parent[i - 1][parent[i - 1][j]];
    }
}

a와 b의 최소 공통 조상을 구하기 위해 먼저 두 정점의 깊이를 구한다.

a의 깊이를 감소시키면서 계산을 진행하기 위해 깊이가 큰 정점을 a에 위치시킨다.

두 정점의 깊이를 동일하게 맞추기 위해 \(2^i\) 를 계산하여 a의 깊이에서 해당 값만큼 빼면서 진행한다.

두 정점의 깊이가 같아졌고 a와 b가 같다면 해당 정점이 최소 공통 조상이 된다.

a와 b가 다르다면 두 정점이 같아질 때까지 두 정점의 깊이를 동시에 올리면서 최소 공통 조상을 구한다.

void LCA(int a, int b) {
    // a와 b의 depth를 저장해놓는다
    int dep_a = depth[a];
    int dep_b = depth[b];

    // a를 감소시키면서 계산 진행할 것
    // depth가 큰 것을 a로 위치시킨다
    if (dep_a < dep_b) {
        int temp = dep_a;
        dep_a = dep_b;
        dep_b = temp;

        temp = a;
        a = b;
        b = temp;
    }

    // a와 b의 depth를 맞춘다
    // a의 depth = 13, b의 depth = 2
    // j = 3 -> 13, 2
    // 2^3 = 8 -> (13 - 8) >= 2 (8만큼 더 올라갈 수 있음)
    // j = 2 -> 5, 2
    // 2^2 = 4 -> (5 - 4) >= 2
    // j = 1 -> 5, 2
    // 2^1 = 2 -> (5 - 2) >= 2 (2만큼 더 올라갈 수 있음)
    // j = 0 -> 3, 2
    // 2^0 = 1 -> (3 - 1) >= 2 (1만큼 더 올라갈 수 있음)
    for (int i = 17; i >= 0; i--) {
        // Math.pow(a, b) -> a^b
        // 2^i 를 가져온다
        // "a의 depth - 2^i >= b의 depth" 보다 클 경우
        if (dep_a - (int) Math.pow(2, i) >= dep_b) {
            // depth 조절
            dep_a -= (int) Math.pow(2, i);
            // a의 i번째 조상
            a = parent[i][a];

            if (dep_a == dep_b)
                break;
        }
    }

    int answer = a;

    // 두 a, b를 같이 올리는 과정
    if (a != b) {
        for (int i = 17; i >= 0; i--) {
            if (parent[i][a] != parent[i][b]) {
                // a의 2^i번째 조상과 b의 2^i번째 조상이 같다면 a와 b를 올린다
                a = parent[i][a];
                b = parent[i][b];
            }

            // 올라가다 보면 LCA에 도달한다
            answer = parent[i][a];
        }
    }

    System.out.println(answer);
}

 

전체 코드

import java.io.*;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

public class Main {

    static int N, M;
    static List<List<Integer>> adjList;
    static int[][] parent;
    static boolean[] check;
    static int[] depth;

    static void getDepth(int cur, int dep) {
        // 한번 방문한 곳은 방문하지 않음
        check[cur] = true;
        // 현재 정점에 depth를 저장한다
        depth[cur] = dep;

        // 인접리스트 (adjList.get(cur).size() -> cur 점의 degree)
        for (int i = 0; i < adjList.get(cur).size(); i++) {
            int dest = adjList.get(cur).get(i);

            // 방문하지 않은 점일 경우
            if (!check[dest]) {
                // dest 점의 직계조상은 cur이다
                parent[0][dest] = cur;
                // DFS 진행
                getDepth(dest, dep + 1);
            }
        }
    }

    static void LCA(int a, int b) {
        // a와 b의 depth를 저장해놓는다
        int dep_a = depth[a];
        int dep_b = depth[b];

        // a를 감소시키면서 계산 진행할 것
        // depth가 큰 것을 a로 위치시킨다
        if (dep_a < dep_b) {
            int temp = dep_a;
            dep_a = dep_b;
            dep_b = temp;

            temp = a;
            a = b;
            b = temp;
        }

        // a와 b의 depth를 맞춘다
        // a의 depth = 13, b의 depth = 2
        // j = 3 -> 13, 2
        // 2^3 = 8 -> (13 - 8) >= 2 (8만큼 더 올라갈 수 있음)
        // j = 2 -> 5, 2
        // 2^2 = 4 -> (5 - 4) >= 2
        // j = 1 -> 5, 2
        // 2^1 = 2 -> (5 - 2) >= 2 (2만큼 더 올라갈 수 있음)
        // j = 0 -> 3, 2
        // 2^0 = 1 -> (3 - 1) >= 2 (1만큼 더 올라갈 수 있음)
        for (int i = 17; i >= 0; i--) {
            // Math.pow(a, b) -> a^b
            // 2^i 를 가져온다
            // "a의 depth - 2^i >= b의 depth" 보다 클 경우
            if (dep_a - (int) Math.pow(2, i) >= dep_b) {
                // depth 조절
                dep_a -= (int) Math.pow(2, i);
                // a의 i번째 조상
                a = parent[i][a];

                if (dep_a == dep_b)
                    break;
            }
        }

        int answer = a;

        // 두 a, b를 같이 올리는 과정
        // depth를 7이라고 가정
        // a != b, dep_a = 7, dep_b = 7
        // 2^3 = 8
        if (a != b) {
            for (int i = 17; i >= 0; i--) {
                if (parent[i][a] != parent[i][b]) {
                    // a의 2^i번째 조상과 b의 2^i번째 조상이 같다면 a와 b를 올린다
                    a = parent[i][a];
                    b = parent[i][b];
                }

                // 올라가다 보면 LCA에 도달한다
                answer = parent[i][a];
            }
        }

        System.out.println(answer);
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());

        check = new boolean[N + 1];
        depth = new int[N + 1];
        parent = new int[18][N + 1];

        adjList = new ArrayList<>();
        for (int i = 0; i <= N; i++) {
            adjList.add(new ArrayList<>());
        }

        for (int i = 1; i < N; i++) {
            st = new StringTokenizer(br.readLine());
            int v1 = Integer.parseInt(st.nextToken());
            int v2 = Integer.parseInt(st.nextToken());

            adjList.get(v1).add(v2);
            adjList.get(v2).add(v1);
        }
        
        // depth를 구하는 탐색
        getDepth(1, 0);

        // 2^17까지 조상을 구해놓는다면 배열 값이 부족할 일이 없음 (2^17 > N = 100,000)
        for (int i = 1; i <= 17; i++) {
            for (int j = 1; j <= N; j++) {
                // j의 2^i번째 부모는 j의 2^(i-1)번째 부모의 2^(i-1)번째 부모이다
                parent[i][j] = parent[i - 1][parent[i - 1][j]];
            }
        }

        st = new StringTokenizer(br.readLine());
        M = Integer.parseInt(st.nextToken());
        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            LCA(a, b);
        }
    }
}