[BOJ] 백준 15480번 LCA와 쿼리 (Java)

#15480 LCA와 쿼리

난이도 : 플레 2

유형 : 트리 / LCA 

15480번: LCA와 쿼리

▸ 문제

N개의 정점으로 이루어져 있는 트리 T가 주어졌을 때, 다음 쿼리를 수행하는 프로그램을 작성하시오.

• r u v: T의 루트가 r이라고 했을 때, u와 v의 LCA를 출력한다.

▸ 입력

첫째 줄에 정점의 개수 N(1 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N-1개의 줄에는 트리 T의 간선 정보 u와 v가 주어지다. u와 v는 트리의 간선을 나타내는 두 정점이다.

다음 줄에는 쿼리의 개수 M(1 ≤ M ≤ 100,000)가 주어진다. 다음 M개의 줄에는 쿼리를 나타내는 r, u, v가 주어진다.

▸ 출력

각각의 쿼리마다 한 줄에 하나씩 결과를 출력한다.

 

문제 풀이  

기존 LCA 문제에서 쿼리 부분이 업그레이드된 문제이다. u, v의 LCA를 구해야 하는데 루트가 고정되어 있지 않아 n개의 트리를 가정하고 구해줘야 한다. 그런데 최대 10만개 정점을 가지는 트리를 각 루트를 정해서 구하는 것은 시간이 너무 걸린다. 그래서 쿼리 연산 부분에서 단축시켜주는 방법을 고려해야 한다.

 

그 방법으로는 LCA(r,a), LCA(r,b), LCA(a,b)를 모두 구하여 가장 depth가 깊은 정점을 선택해주면 된다. 사실 증명이랄 것은 아니지만 3개의 정점이 트리에 존재할 수 있는 위치를 모두 고려해본 다음 정리해보면 다음과 같이 위의 방식이 옳음을 알 수 있다.

 

LCA구하기

 

💡 LCA에 대한 자세한 설명은 여기를 참고해주세요

 

설계

1. 양방향 트리 데이터를 입력받는다.
2. 1번 정점을 루트로하고 LCA를 구하기 위한 데이터를 초기화시켜준다.
   
   1. depth[cur]: 해당 정점(cur)의 높이
   2. parent[cur][h]: 해당 정점(cur)의 2^h번째 부모
3. (r, u, v)의 LCA를 구하기 위해서는 위에서 구한 1번 정점이 루트인 트리를 기준으로 LCA(r, a), LCA(r, b), LCA(a, b) 중 가장 depth가 깊은 LCA를 선택해주면 된다.
   1. sb.append(comparingDepth(LCA(r,a), comparingDepth(LCA(r,b), LCA(a,b)))+"\n");
4. 그렇게 구한 r을 루트로하는 u, v의 LCA를 출력해준다.

 

풀이 코드 

import java.io.*;
import java.util.*;

public class Main {

	static int n, h;
	static List<Integer>[] list;
	static int[] depth;
	static int[][] parent;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		
		n = Integer.parseInt(br.readLine());
		list = new ArrayList[n+1];
		for(int i=1; i<n+1; i++){
			list[i] = new ArrayList<>();
		}
		
		StringTokenizer st;
		for(int i=0; i<n-1; i++) {
			st = new StringTokenizer(br.readLine());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			list[a].add(b);
			list[b].add(a);
		}
	
		h = (int)Math.ceil(Math.log(n)/Math.log(2)) +1;
		parent = new int[n+1][h];
		depth = new int[n+1];
		
		initTree(1,0);
		fillParents();
		
		StringBuilder sb = new StringBuilder();
		int m = Integer.parseInt(br.readLine());
		for(int i=0; i<m; i++) {
			st = new StringTokenizer(br.readLine());
			int r = Integer.parseInt(st.nextToken());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			sb.append(comparingDepth(LCA(r,a), comparingDepth(LCA(r,b), LCA(a,b)))+"\n");
		}
		System.out.println(sb.toString());
	}
	
	static int comparingDepth(int a, int b) {
		if(depth[a] > depth[b]) {
			return a;
		}else {
			return b;
		}
	}
	
	static int LCA(int a, int b) {
		int ah = depth[a];
		int bh = depth[b];
		if(ah < bh) {
			int tmp = a;
			a = b;
			b = tmp;
		}
		
		for(int i=h-1; i>=0; i--) {
			if(Math.pow(2, i) <= depth[a] - depth[b]) {
				a = parent[a][i];
			}
		}
		
		if(a == b) return a;
		
		for(int i=h-1; i>=0; i--) {
			if(parent[a][i] != parent[b][i]) {
				a = parent[a][i];
				b = parent[b][i];
			}
		}
		return parent[a][0];
	}
	
	static void initTree(int idx, int pa) {
		for(int nxt : list[idx]) {
			if(nxt != pa) {
				depth[nxt] = depth[idx]+1;
				initTree(nxt, idx);
				parent[nxt][0] = idx; 
			}
		}
	}
	
	static void fillParents() {
		for(int i=1; i<h; i++) {
			for(int j=1; j<n+1; j++) {
				parent[j][i] = parent[parent[j][i-1]][i-1];
			}
		}
	}
}