백준 2233 사과나무

문제

주어진 이진수열로 트리를 구성하고, 썩은 사과의 정보가 주어졌을 때 하나의 사과만을 잘랐을 때 가장 적은 사과를 잃는 하나의 사과를 구하는 문제
자세한 내용

입력

첫째 줄에 정점의 개수 N(1≤N≤2,000)이 주어진다. 둘째 줄에는 벌레가 만드는 2×N자리의 이진수가 주어진다. 셋째 줄에는 썩은 사과의 위치를 나타내는 두 정수 X, Y가 주어진다. 이는 2×N자리의 이진수에서 X번째의 숫자에 대응되는 정점과, Y번째 숫자에 대응되는 정점에 있는 사과가 썩었다는 의미이다. 이때 두 정점이 서로 같을 수도 있다.

출력

첫째 줄에 제거해야 할 사과를 나타내는 두 정수 i, j를 출력한다. 제거해야 할 사과를 Z라고 했을 때, 이는 Z를 방문할 때의 0의 위치와 Z에서 리턴될 때의 1의 위치가 이진수에서 각각 i, j 번째임을 나타낸다.

풀이

이진 수열로 트리를 구성하고 LCA

이런 식의 트리 구성이 입력으로 주어지는 건 처음이었지만 문제대로 따라가면 무리는 없었고 LCA를 복기하는 의미가 있었다.

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

public class Main {
    static int N, S;
    static int[] depth, nodeInSeq, zeroInSeq, oneInSeq;
    static int[][] parent; // j번쨰 노드의 2^i 번쨰 부모 : 2^i 번째 부모는 2^(i-1)번쨰 부모의 2^(i-1)번쨰 부모

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringBuilder sb = new StringBuilder();
        StringTokenizer st;

        N = Integer.parseInt(br.readLine());
        S = 0;
        for (int i = 1; i <= N; i *= 2) {
            S++;
        }
        depth = new int[N + 1];
        parent = new int[S][N + 1];
        String seq = br.readLine();
        nodeInSeq = new int[2 * N + 1];
        zeroInSeq = new int[N + 1];
        oneInSeq = new int[N + 1];
        int len = seq.length();
        int current = 1;
        int id = 1;

        for (int i = 1; i <= len; i++) {
            if (seq.charAt(i - 1) == '0') {
                parent[0][id] = current;
                depth[id] = depth[current] + 1;
                nodeInSeq[i] = id;
                zeroInSeq[id] = i;
                current = id++;
            } else {
                nodeInSeq[i] = current;
                oneInSeq[current] = i;
                current = parent[0][current];
            }
        }
        for (int i = 1; i < S; i++) {
            for (int j = 1; j <= N; j++) {
                parent[i][j] = parent[i - 1][parent[i - 1][j]];
            }
        }
        st = new StringTokenizer(br.readLine());
        int a = Integer.parseInt(st.nextToken());
        int b = Integer.parseInt(st.nextToken());
        int A = nodeInSeq[a];
        int B = nodeInSeq[b];
        int lca = lca(A, B);
        bw.write(zeroInSeq[lca] + " " + oneInSeq[lca] + "\n");
        bw.flush();
        bw.close();
    }
    static int lca(int a, int b) {
        if(depth[a] > depth[b]) { // 항상 b가 더 깊도록
            int temp = a;
            a = b;
            b = temp;
        }

        for (int i = S - 1; i >= 0; i--) { // 깊이 맞춰주기
            if (depth[a] <= depth[parent[i][b]]) {
                b = parent[i][b];
            }
        }
        if(a == b) return a;

        for (int i = S - 1; i >= 0; i--) {
            if (parent[i][a] != parent[i][b]) {
                a = parent[i][a];
                b = parent[i][b];
            }
        }
        return parent[0][a];
    }
}