SWEA 1855번: 영준이의 진짜 BFS Java Solution

1. Problem Definition & Analysis

주어진 트리에서 1번 정점을 시작으로 하여 BFS를 수행했을 때의 방문 순서가 있을 것이다. 이 방문 순서에 따라 $i$번째로 방문되는 정점에서 $i + 1$번째로 방문되는 정점으로 트리를 따라 이동한다. 이때 총 몇 번의 이동이 발생하는지 집계하는 것을 문제에서 요구하고 있다.

문제를 해결하는 대략적인 프로세스는 다음과 같다.

1. BFS를 통해 정점의 방문 순서를 구한다.
2. 정점의 방문 순서에서 인접한 모든 정점 쌍에 대해 다음과 같은 반복을 수행한다.
   1. 출발 정점($v1$), 도착 정점($v2$) 쌍에 대해 LCA 알고리즘으로 공통 조상($p$)을 구한다.
   2. 어떤 정점 $v$의 깊이를 $d_{v}$라고 하자. $(d_{v1} - d_{p}) + (d_{v2} - d_{p}) = d_{v1} + d_{v2} - 2 \times d_{p}$와 같은 식을 통해 $v1$에서 $v2$의 이동 횟수를 계산한다.

이때 트리의 깊이가 $N$인 경우를 생각하면 전체 이동 횟수가 4 바이트 정수형의 범위를 초과할 수 있다는 것에 주의해야 한다.

2. Solution

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

public class Solution {

    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static StringTokenizer st;
    static StringBuilder sb = new StringBuilder();

    static int T, N;
    static int[] depth;
    static int[][] parent;
    static List<List<Integer>> tree;

    public static void main(String[] args) throws IOException {
        T = Integer.parseInt(br.readLine());
        for (int i = 1; i <= T; i++)
            sb.append('#').append(i).append(' ').append(solve()).append('\n');
        System.out.println(sb);
    }

    public static long solve() throws IOException {
        initialize();
        long totalDistance = 0;
        List<Integer> visitedOrder = BFS();
        for (int i = 1; i < visitedOrder.size(); i++)
            totalDistance += getDistance(visitedOrder.get(i), visitedOrder.get(i - 1));
        return totalDistance;
    }

    public static void initialize() throws IOException {
        N = Integer.parseInt(br.readLine());

        parent = new int[N + 1][(int) Math.ceil(Math.log(N) / Math.log(2)) + 1];
        depth = new int[N + 1];
        tree = new ArrayList<>();
        for (int i = 0; i <= N; i++)
            tree.add(new ArrayList<>());

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

        for (int pow = 1; pow < parent[0].length; pow++)
            for (int c = 1; c <= N; c++)
                parent[c][pow] = parent[parent[c][pow - 1]][pow - 1];
    }

    public static int getDistance(int v1, int v2) {
        int lca = LCA(v1, v2);
        return depth[v1] + depth[v2] - 2 * depth[lca];
    }

    public static int LCA(int v1, int v2) {
        if (depth[v1] < depth[v2]) { int t = v1; v1 = v2; v2 = t; }

        while (depth[v1] != depth[v2]) {
            int diff = depth[v1] - depth[v2];
            int jump = (int) (Math.log(diff) / Math.log(2));
            v1 = parent[v1][jump];
        }

        while (v1 != v2) {
            int k = (int) (Math.log(depth[v1]) / Math.log(2));
            while (k > 0 && parent[v1][k] == parent[v2][k])
                k--;
            v1 = parent[v1][k];
            v2 = parent[v2][k];
        }

        return v1;
    }

    public static List<Integer> BFS() {
        List<Integer> sequence = new ArrayList<>();
        Queue<Integer> queue = new ArrayDeque<>();
        queue.offer(1);
        sequence.add(1);
        while (!queue.isEmpty()) {
            int curr = queue.poll();
            for (int child : tree.get(curr)) {
                queue.offer(child);
                sequence.add(child);
            }
        }
        return sequence;
    }
}

BFS, LCA 로직은 정석적이고 기초 문제가 아니기 때문에 해당 알고리즘에 대한 자세한 설명은 생략하겠다.