[1855. 영준이의 진짜 BFS]

문제


노드1에서 노드2로 이동할 때 결국 공통 조상을 통과한다. 노드1에서 공통 조상까지는 depth[노드1] - depth[공통조상], 공통조상에서 노드2까지의 거리는 depth[노드2] - depth[공통조상] 이다. 결국 depth[노드1] + depth[노드2] - depth[공통조상] - depth[공통조상] 가 이동 경로인 것을 알 수 있다. 공통 조상의 depth 를 구하기 위해서 LCA 알고리즘을 사용하였다.

Queue 로 BFS 를 하면서 child 를 정렬해서 방문한다. 이때 현재 노드에서 다음 노드로의 이동 경로를 계속해서 합해준다.

포인트


  1. List[] childs = new ArrayList[N+1]; list의 배열을 만드는 방법이다. 이때 for 문으로 돌면서 모두 초기화를 해줘야한다.
  2. sum 을 int 로 초기화하니 계속해서 오답이 발생했다. int 값을 초과해 long 으로 바꿔줘야 정답으로 처리된다. 시간과 반복 횟수에 주의해서 타입을 설정하자.
  3. 단순하게 번호순서대로 이동한다고 생각했다. 그것이 아니라 BFS 내부에서 가장 작은 자식부터 방문하는 것이다. 문제를 꼼꼼하게 잘 읽자!

코드


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

public class Solution {
	static StringBuilder result = new StringBuilder();
	static int[] depth = new int[100001];
	static int[][] parents = new int[100001][17];


	public static void main(String args[]) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		int T = Integer.parseInt(br.readLine());

		///////////////////////////
		int N;
		///////////////////////////
		for (int test_case = 1; test_case <= T; test_case++) {
			result.append("#" + test_case + " ");
			/////////////////////////////
			Queue<Integer> queue = new LinkedList();
			N = Integer.parseInt(br.readLine());

			List<Integer>[] childs = new ArrayList[N + 1];
			for (int i = 0; i < N; i++) {
				childs[i] = new ArrayList<Integer>();
			}

			//log2 N 구하기..
			int k = 1;
			int K = 0;
			while (k < N) {
				k <<= 1;
				K++;
			}

			depth[1] = 1;

			StringTokenizer st = new StringTokenizer(br.readLine());
			for (int i = 2; i <= N; i++) {
				int p = Integer.parseInt(st.nextToken());
				depth[i] = depth[p] + 1;
				parents[i][0] = p;
				childs[p].add(i);
			}

			//parent 배열 완성
			for (int i = 1; i < K; i++) {
				for (int j = 1; j <= N; j++) {
					parents[j][i] = parents[parents[j][i - 1]][i - 1];
				}
			}

			long sum = 0;
			queue.add(1);
			int idx = 1;
			while (!queue.isEmpty()) {
				sum = sum + depth[idx] + depth[queue.peek()] - (2 * depth[lca(idx, queue.peek(), K)]);
				idx = queue.poll();
				if (childs[idx] != null) {
					childs[idx].sort(Comparator.naturalOrder());
					for (int c :
							childs[idx]) {
						queue.add(c);
					}
				}
			}

			result.append(sum + "\n");
			System.out.print(result);
			result.setLength(0);

			for (int i = 1; i <= N; i++) {
				depth[i] = 0;
				for (int j = 0; j < K; j++) {
					parents[i][j] = 0;
				}
			}
		}
	}

	static int lca(int a, int b, int K) {
		if (depth[a] < depth[b]) {
			int temp = a;
			a = b;
			b = temp;
		}

		for (int i = K - 1; i >= 0; i--) {
			if (Math.pow(2, i) <= depth[a] - depth[b]) {
				a = parents[a][i];
			}
		}

		if (a == b) return a;

		for (int i = K - 1; i >= 0; i--) {
			if (parents[a][i] != parents[b][i]) {
				a = parents[a][i];
				b = parents[b][i];
			}
		}

		return parents[a][0];
	}
}

업데이트:

댓글남기기