SWEA 1248 - 공통조상

문제 링크

문제 설명

이진 트리에서 임의의 두 정점의 가장 가까운 공통 조상을 찾고, 그 정점을 루트로 하는 서브 트리의 크기 구하기

접근

• 이진 트리이기에 트리를 배열로 저장할 수 있음
• 공통 조상 찾기 -> Union find 사용
  Union Find 정리 바로가기
• 서브 트리의 크기 구하기 -> dfs 사용해 구현

정답 코드 (자바/java)

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.Arrays;

public class SWEA_1248 {
	static int[][] graph;

	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int testCase = Integer.parseInt(br.readLine()) + 1;
		int v, e, v1, v2, size, par, child, tmp1, tmp2, cnt;
		String[] str;
		StringBuilder sb = new StringBuilder();

		for (int tc = 1; tc < testCase; tc++) {
			sb.append('#').append(tc).append(" ");

			str = br.readLine().split(" ");
			// 정점, 간선, 정점1, 정점2 입력받기
			v = Integer.parseInt(str[0]);
			e = Integer.parseInt(str[1]);
			v1 = Integer.parseInt(str[2]);
			v2 = Integer.parseInt(str[3]);

			// 간선을 저장할 그래프
			graph = new int[v + 1][2];
			// 부모를 저장할 배열 -> union find를 위해서
			int[] parent = new int[v + 1];

			str = br.readLine().split(" ");
			size = str.length;

			// 그래프 저장
			for (int i = 0; i < size; i = i + 2) {
				par = Integer.parseInt(str[i]);
				child = Integer.parseInt(str[i + 1]);

				if (graph[par][0] == 0)
					graph[par][0] = child;
				else
					graph[par][1] = child;

				parent[child] = par;
			}

			// union find 사용해 공통 조상 찾기
			// 첫번째 노드의 조상 다 -1로 설정
			tmp1 = v1;
			while (parent[tmp1] != 0) {
				tmp2 = tmp1;
				tmp1 = parent[tmp1];
				parent[tmp2] = -1;
			}
			parent[tmp1] = -1;

			// 두번째 노드의 조상 중 -1로 설정된 것(겹치는 것) 찾기
			tmp1 = v2;
			while (parent[tmp1] != -1) {
				tmp1 = parent[tmp1];
			}

			// 공통 조상 : tmp1 일 때, 하위 노드 수 세기
			cnt = dfs(tmp1, 0);
			sb.append(tmp1).append(" ").append(cnt).append('\n');
		}
		System.out.println(sb);
	}

	// 하위 노드 수 세기
	private static int dfs(int tmp1, int cnt) {
		int tmpCnt = cnt;
		
		// 해당 노드가 마지막일 때, cnt + 1 한 후 출력
		if (graph[tmp1][0] == 0) {
			return cnt + 1;
		}
		
		// 왼쪽 하위 노드 서치
		tmpCnt += dfs(graph[tmp1][0], cnt);
		
		// 오른쪽 하위 노드 없을 때, +1
		if (graph[tmp1][1] == 0) {
			return tmpCnt + 1;
		}
		
		// 오른쪽 하위 노드 있을 때
		tmpCnt += dfs(graph[tmp1][1], cnt);
		return tmpCnt + 1;
	}

}

반성

dfs를 통해 하위 노드의 수를 셀 때, 조금 더 간단한 방법으로 구현이 가능할 것 같음