[백준/Java] 11438 LCA 2

문제

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

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

입력

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

출력

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

---

- 부모 노드 배열 점화식

P[K][N] = P[K - 1][P[K - 1][N]]

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

public class Main {
	static ArrayList<Integer> tree[];    // 인접 리스트 자료구조
	static int depth[];        // 노드 깊이 배열
	static int kmax;           // 최대 가능 높이
	static int parent[][];     // 노드 조상 배열
	static boolean visited[];  // 방문 여부 저장 배열
	
	public static void main(String[] args) throws NumberFormatException, IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		
		int n = Integer.parseInt(br.readLine());    // 노드 개수 입력 받음
		tree = new ArrayList[n + 1];        // 인접 리스트 선언(인덱스 1부터 시작이라 +1)
		for (int i = 1; i <= n; i++) {      // 인접 리스트 초기화, 인덱스 1부터 시작
			tree[i] = new ArrayList<Integer>();
		}
		
		for (int i = 0; i < n - 1; i++) {   // 노드 입력 받아서 데이터에 저장
			StringTokenizer st = new StringTokenizer(br.readLine());
			int s = Integer.parseInt(st.nextToken());
			int e = Integer.parseInt(st.nextToken());
			tree[s].add(e);    // 인접 리스트니까 연결
			tree[e].add(s);
		}
		
        
        // 인덱스 1부터 시작하는 걸로 할거라 배열 크기는 n + 1
		depth = new int[n+1];            // 각 노드 별 깊이를 나타내는 배열
		visited = new boolean[n+1];      // 각 노드 별 방문했는지 나타내는 배열
		int temp = 1;
		kmax = 0;
		
		while (temp <= n) {    // 최대 가능 높이 구하기(kmax) ex) n: 15 이면 kmax: 4
			temp <<= 1;        // temp = temp << 1 (비트 연산) -----> 1 2 4 8 16 32 ...
			kmax++;
		}
		
        
		parent = new int[kmax + 1][n + 1];    // 크기 선언, 2^k까지의 부모 노드 저장
		BFS(1);        // depth와 부모 노드를 BFS를 이용해 구하기
		
		for (int k = 1; k <= kmax; k++) {    // kmax만큼 반복
			for (int j = 1; j <= n; j++) {   // 노드 개수만큼 반복
				parent[k][j] = parent[k - 1][parent[k - 1][j]];    // 부모 노드 배열 점화식 이용
			}
		}
		
        
		int m = Integer.parseInt(br.readLine());    // 질의 개수
		for (int i = 0; i < m; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
            // 최소 공통 조상을 구할 두 노드 a, b
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			int LCA = excuteLCA(a, b);
			System.out.println(LCA);
		}
	}
	
    // LCA를 구하는 함수
	static int excuteLCA(int a, int b) {
		if (depth[a] > depth[b]) {    // 각 노드 깊이 비교, 더 깊은 depth가 b가 되도록 변경하기(값이 더 크다는건 더 깊음)
			int temp = a;    // 조건 맞으면 노드 swap
			a = b;
			b = temp;
		}
		
		for (int k = kmax; k >= 0; k--) {       // depth 맞추기
			if (Math.pow(2, k) <= depth[b] - depth[a]) {    // parent 배열을 이용해 2의 제곱수로 이동
				if (depth[a] <= depth[parent[k][b]]) {
					b = parent[k][b];
				}
			}
		}
		
		for (int k = kmax; k >= 0; k--) {      // 조상 찾기
			if (parent[k][a] != parent[k][b]) {
				a = parent[k][a];
				b = parent[k][b];
			}
		}
		
		int LCA = a;
		if (a != b) LCA = parent[0][LCA];
		
		return LCA;    // 최소 공통 조상 리턴
	}
	
    
    // BFS : 너비 우선 탐색 <-> DFS
	private static void BFS(int node) {
		Queue<Integer> queue = new LinkedList<Integer>();
		queue.add(node);        // queue 자료구조에 출발 노드 더하기
		visited[node] = true;   // visited 배열에 현재 노드 방문 기록
		int level = 1;
		int now_size = 1;
		int count = 0;
		
		while (!queue.isEmpty()) {    // 큐가 empty일 때까지 반복
			int now_node = queue.poll();    // 큐에서 노드 데이터 가져오기(poll)
			for (int next : tree[now_node]) {    // 현재 노드의 연결 노드 중 방문하지 않은 노드로 반복
				if (!visited[next]) {    // 큐에 데이터 삽입하고, visited 배열에 방문 기록
					visited[next] = true;
					queue.add(next);  
					parent[0][next] = now_node;    // parent 배열에 자신의 부모 노드 저장
					depth[next] = level;           // depth 배열에 현재 높이 저장하기
				}
			}
			count++;
			
			if (count == now_size) {    // 이번 높이에 해당하는 모든 노드를 방문했을 때
				count = 0;
				now_size = queue.size();
				level++;
			}
		}
	}
}