[백준] 11438 LCA 2 - 자바

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

11438번: LCA 2

문제 풀이

데이터가 적은 기본적인 LCA 문제는 a, b 각 노드에 대해 부모 노드를 찾아 올라가며 공통 조상 노드를 찾아야 한다.

이 과정을 줄여 각 깊이별 부모 노드를 2차원 배열을 통해 관리하여 더욱 효율적으로 구현할 수 있는데 이 문제가 해당 알고리즘을 적용하여야 풀 수 있는 문제이다.

 

주어진 예제에 대해 어떻게 풀어나갈지 흐름대로 살펴보자

예제 트리

트리 좌측은 각 노드의 레벨 (깊이)을 나타낸다. 기본적인 LCA 알고리즘에서는 각 노드의 부모노드만 알고 순차탐색 했다면 이 문제에서 적용할 LCA 문제에서는 각 노드의 레벨별 부모 즉, 조상들까지 dp로 관리해야 한다. 

dp[i][j]: i노드의 j+1번째 조상의 노드, 즉 dp[15][1]: 15번 노드의 2번째 조상이므로 5가 된다. 따라서 dp 크기는 트리의 높이에 의해 정해진다. 트리의 높이는 (logN/log2)이고 소수점이 생길 수 있으니 올림을 해줘서 구한다.

 

dp의 크기를 정했다면 표를 채워보자. 이렇게 모두 구해놓으면 LCA를 구하기 위해 두 노드값이 주어졌을 때 매번 부모노드를 타고 올라가는 것보다 훨씬 효율적일 것이다.

parents[i][j]	0	1	2	3
1 (루트)	0	0	0	0
2	1	0	0	0
3	1	0	0	0
4	2	1	0	0
5	2	1	0	0
6	2	1	0	0
7	3	1	0	0
8	3	1	0	0
9	4	2	1	0
10	4	2	1	0
11	5	2	1	0
12	5	2	1	0
13	7	3	1	0
14	7	3	1	0
15	11	5	1	0

이제 두 노드가 주어졌을 때 같은 레벨로 맞춘 후 순차적으로 두 노드의 부모값이 같아질 때까지 부모 노드로 승격시키며 공통 조상 노드를 찾을 수 있다.

 

이제 코드로 살펴보자.

먼저 LCA를 실행시키기 위해 초기값을 세팅해줘야 하기 때문에 tree 관계도 저장해야 한다. 두 노드의 레벨(depth)도 필요하기 때문에 추가로 할당해준다.

// init
tree = new ArrayList[N + 1];
for (int i = 1; i <= N; i++) {
    tree[i] = new ArrayList<>();
}
heightOfTree = (int) ceil(log(N) / log(2));
depths = new int[N + 1];
parents = new int[N + 1][heightOfTree];

tree의 관계를 설정해준다.

for (int i = 0; i < N-1; i++) {
    st = new StringTokenizer(br.readLine());
    int from = Integer.parseInt(st.nextToken());
    int to = Integer.parseInt(st.nextToken());

    tree[from].add(to);
    tree[to].add(from);
}

 

재귀함수를 통해 각 노드의 depth와 바로 위 부모 노드 parents[][0]를 설정해준다. 루트->리프 노드로 방문하고 있고 부모 노드로 재방문하는 것을 막기 위해 조건문을 추가해주었다.

// 루트는 1번이므로 1번부터 depths[]와 해당 parents[][0]: 해당 노드의 바로 윗 단계 parent 부모 설정
initRecur(1, 1, 0);

static void initRecur(int idx, int depth, int parent) {
    depths[idx] = depth;
    for (int adj : tree[idx]) {
        // 다시 부모로 방문하는 것을 방지, visited 역할
        if (adj != parent) {
            initRecur(adj, depth + 1, idx);
            parents[adj][0] = idx;
        }
    }
}

위에서 설정한 parents[][0]을 활용하여 나머지 부분을 채워주자. 기존에 저장된 부모 노드를 따라가며 채우는 코드이다. 이해가 힘들다면 위에서 미리 채워둔 표를 활용하여 따라가보자.

// 부모의 부모를 탐색하며 parents[][1~h]: 조상 노드값 채우기
static void fillAncestor() {
    for (int h = 1; h < heightOfTree; h++) {
        for (int idx = 1; idx < N + 1; idx++) {
            parents[idx][h] = parents[parents[idx][h - 1]][h - 1];
        }
    }
}

이제 마지막으로 두 노드값을 입력받아서 같은 레벨로 맞추고 공통 조상 노드를 찾는다.

static int LCA(int a, int b) {
    int ah = depths[a];
    int bh = depths[b];

    // ah > bh로 세팅
    if(ah < bh) {
        int tmp = a;
        a = b;
        b = tmp;
    }

    // a가 b보다 아래에 있다면 b와 같은 깊이로 올려줌
    for (int h = heightOfTree; h >= 0; h--) {
        if (pow(2, h) <= depths[a] - depths[b]) { // 포화이진트리의 성질
            a = parents[a][h];
        }
    }

    if (a == b) {
        return a; // or b
    }

    // a와 b의 부모가 같아질 때까지 a와 b를 부모 노드로 승격
    for (int h = heightOfTree - 1; h >= 0; h--) {
        if (parents[a][h] != parents[b][h]) {
            a = parents[a][h];
            b = parents[b][h];
        }
    }

    return parents[a][0]; // or parent[b][0]
}

코드

package 백준;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

import static java.lang.Math.*;

public class b11438 {

    static int N, M;
    static List<Integer>[] tree;
    static int[] depths;
    static int[][] parents;
    static int heightOfTree;

    public static void main(String[] args) throws IOException {
        StringBuilder sb = new StringBuilder();
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
        N = Integer.parseInt(br.readLine());

        // init
        tree = new ArrayList[N + 1];
        for (int i = 1; i <= N; i++) {
            tree[i] = new ArrayList<>();
        }
        heightOfTree = (int) ceil(log(N) / log(2));
        depths = new int[N + 1];
        parents = new int[N + 1][heightOfTree];

        for (int i = 0; i < N-1; i++) {
            st = new StringTokenizer(br.readLine());
            int from = Integer.parseInt(st.nextToken());
            int to = Integer.parseInt(st.nextToken());

            tree[from].add(to);
            tree[to].add(from);
        }

        // 루트는 1번이므로 1번부터 depths[]와 해당 parents[][0]: 해당 노드의 바로 윗 단계 parent 부모 설정
        initRecur(1, 1, 0);
        fillAncestor();

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

            sb.append(LCA(a, b) + "\n");
        }

        System.out.println(sb);
    }

    static void initRecur(int idx, int depth, int parent) {
        depths[idx] = depth;
        for (int adj : tree[idx]) {
            // 다시 부모로 방문하는 것을 방지, visited 역할
            if (adj != parent) {
                initRecur(adj, depth + 1, idx);
                parents[adj][0] = idx;
            }
        }
    }

    // 부모의 부모를 탐색하며 parents[][1~h]: 조상 노드값 채우기
    static void fillAncestor() {
        for (int h = 1; h < heightOfTree; h++) {
            for (int idx = 1; idx < N + 1; idx++) {
                parents[idx][h] = parents[parents[idx][h - 1]][h - 1];
            }
        }
    }

    static int LCA(int a, int b) {
        int ah = depths[a];
        int bh = depths[b];

        // ah > bh로 세팅
        if(ah < bh) {
            int tmp = a;
            a = b;
            b = tmp;
        }

        // a가 b보다 아래에 있다면 b와 같은 깊이로 올려줌
        for (int h = heightOfTree; h >= 0; h--) {
            if (pow(2, h) <= depths[a] - depths[b]) { // 트리의 성질
                a = parents[a][h];
            }
        }

        if (a == b) {
            return a; // or b
        }

        // a와 b의 부모가 같아질 때까지 a와 b를 부모 노드로 승격
        for (int h = heightOfTree - 1; h >= 0; h--) {
            if (parents[a][h] != parents[b][h]) {
                a = parents[a][h];
                b = parents[b][h];
            }
        }

        return parents[a][0]; // or parent[b][0]
    }
}