[BOJ] 11438번 LCA2

개요

트리의 최소 공통 조상을 찾는 문제

LCA문제보다 입력의 크기가 커져 알고리즘을 최적화 해야 한다.

문제링크

코드

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

public class Baekjoon11438 {
    public static ArrayList<Integer>[] tree;
    public static int[] depths;
    public static int[][] parents;
    public static boolean[] isVisted;
    public static final int MAX_D = 17;
    public static void main(String[] args) throws IOException{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st;

        int n = Integer.parseInt(br.readLine());
        tree = new ArrayList[n+1];
        depths = new int[n+1];
        parents = new int[n+1][MAX_D+1];
        isVisted = new boolean[n+1];
        for (int i = 0; i < tree.length; i++) 
            tree[i] = new ArrayList<>();
        
        int from, to;
        for (int i = 0; i < n-1; i++) {
            st = new StringTokenizer(br.readLine());
            from = Integer.parseInt(st.nextToken());
            to = Integer.parseInt(st.nextToken());

            tree[from].add(to);
            tree[to].add(from);
        }
        isVisted[1] = true;
        dfs(1, 0);

        for (int i = 0; i < MAX_D-1; i++) {
            for (int j = 1; j < n+1; j++) {
                if(parents[j][i] != 0)  //부모배열 채우기
                    parents[j][i+1] = parents[parents[j][i]][i];
            }
        }
        int m = Integer.parseInt(br.readLine());
        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            from = Integer.parseInt(st.nextToken());
            to = Integer.parseInt(st.nextToken());

            bw.write(LCA(from, to)+"\n");
        }
        bw.flush();

    }

    public static void dfs(int node, int depth) {
        int child;
        depths[node] = depth;
        for (int i = 0; i < tree[node].size(); i++) {
            child = tree[node].get(i);
            if(!isVisted[child]) {
                isVisted[child] = true;
                parents[child][0] = node;
                dfs(child, depth + 1);
            }
        }
    }

    public static int LCA(int node1, int node2) {
        if(depths[node1] < depths[node2]) { //node1의 깊이가 더 깊게 변경
            int temp = node1;
            node1 = node2;
            node2 = temp;
        }
        
        for (int i = MAX_D; i >= 0; i--) {  //두 노드의 깊이차이가 2의i승보다 크거나 같으면 2의i승에 해당하는 부모노드로 이동
            if(depths[node1]-depths[node2] >= (1 << i)) {
                node1 = parents[node1][i];
            }
        }
        if(node1 == node2)
            return node1;

        for (int i = MAX_D; i >= 0; i--) {  //두 노드가 같은 깊이이므로 같은 노드를 만날 때까지 부모노드로 2의i승씩 이동
            if(parents[node1][i] != parents[node2][i]) {
                node1 = parents[node1][i];
                node2 = parents[node2][i];
            }
        }
        return parents[node1][0];
    }
}

해결 과정

1. 트리를 탐색하면서 parents배열에 각 노드들의 첫번째 부모노드(parents[i][0])를 저장하고 각 노드의 깊이를 depths배열에 저장한다.

2. 탐색이 끝나면 부모노드배열의 첫번째 열을 가지고 나머지 열을 계산해 부모노드배열을 채운다.

3. 최소 공통 노드를 확인할 노드 2개를 입력받아 두 노드의 깊이가 같아지도록 한 쪽 노드를 parents배열을 이용해 2의 i승만큼 부모노드로 이동하며 깊이를 맞춘다.

4. 깊이가 같아졌으면 노드의 값을 비교해 다르면 2의 i승만큼 부모노드로 이동해 다시 비교하고 노드의 값이 같으면 공통 조상 노드이므로 노드의 값을 반환한다.

이전 문제와 다른 점은 부모노드로 타고 올라갈 때 한칸씩 올라가는 것이 아닌 2의 i승만큼 올라간다는 것이다. 이러면 시간복잡도가O(N)에서 O(logN)정도로 줄어든다.

결과