98. LCA (백준 11437) [JAVA]

문제

문제 링크 : https://www.acmicpc.net/problem/11437

11437번: LCA

문제

N(2 ≤ N ≤ 50,000)개의 정점으로 이루어진 트리가 주어진다. 트리의 각 정점은 1번부터 N번까지 번호가 매겨져 있으며, 루트는 1번이다.

두 노드의 쌍 M(1 ≤ M ≤ 10,000)개가 주어졌을 때, 두 노드의 가장 가까운 공통 조상이 몇 번인지 출력한다.

입력

첫째 줄에 노드의 개수 N이 주어지고, 다음 N-1개 줄에는 트리 상에서 연결된 두 정점이 주어진다. 그 다음 줄에는 가장 가까운 공통 조상을 알고싶은 쌍의 개수 M이 주어지고, 다음 M개 줄에는 정점 쌍이 주어진다.

출력

M개의 줄에 차례대로 입력받은 두 정점의 가장 가까운 공통 조상을 출력한다.

풀이

입력받는 노드들의 부모-자식, 깊이를 알 수 없기 때문에 트리 형태로 정해줘야 한다.

따라서 BFS를 활용해서 depth랑 parent를 정한 다음 LCA를 구하면 된다.

공통 조상을 구하는 방법은, 이전 문제에선 한 노드의 부모를 쭉 구하고 거기서 찾는 방법을 사용했지만

두 노드의 depth를 맞춘 다음 부모로 하나씩 거슬러 올라가는 방법이 더 효율적이다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {

    //root 1
    static int N;
    static int M;
    static Node[] nodes;
    static ArrayList<Integer> answer = new ArrayList<>();

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

        for(int i=0; i<N+1; i++){
            nodes[i] = new Node(99999);
        }
        //root
        nodes[1].depth=0;

        int i1, i2;
        for(int i=1; i<N; i++){
            st = new StringTokenizer(br.readLine());
            i1 = Integer.parseInt(st.nextToken());
            i2 = Integer.parseInt(st.nextToken());
            nodes[i1].linkedNodes.add(i2);
            nodes[i2].linkedNodes.add(i1);
        }

        BFS();

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

            answer.add(LCA(i1, i2));
        }

        for (int ans: answer) {
            System.out.println(ans);
        }
    }

    //루트부터 BFS를 해서 각 노드들의 depth, 부모 결정 가능
    public static void BFS(){
        Queue<Integer> q = new LinkedList<>();
        q.add(1);

        while(!q.isEmpty()){
            int curNum = q.poll();
            Node cur = nodes[curNum];

            for(int childs : cur.linkedNodes){
                if(nodes[childs].depth > cur.depth){
                    nodes[childs].depth = cur.depth+1;
                    nodes[childs].parent = curNum;
                    q.add(childs);
                }
            }
        }
    }

    public static int LCA(int a, int b){
        int aDepth = nodes[a].depth;
        int bDepth = nodes[b].depth;

        if(aDepth != bDepth){
            while(aDepth > bDepth){
                a = nodes[a].parent;
                aDepth--;
            }

            while(aDepth < bDepth){
                b = nodes[b].parent;
                bDepth--;
            }
        }

        while(a != b){
            a = nodes[a].parent;
            b = nodes[b].parent;
        }
        return a;
    }
}

class Node{
    int depth;
    int parent;
    ArrayList<Integer> linkedNodes = new ArrayList<>();

    public Node(int depth){
        this.depth = depth;
    }
}

결과