트리 경로 연결 여부와 에지 제거 문제

문제 설명

주어진 트리 구조에서 특정 두 정점을 연결하는 경로가 존재하는지를 확인하고, 필요에 따라 특정 에지를 제거하는 작업을 수행해야 한다. 트리는 루트가 (1)인 연결 그래프이고, 에지를 제거하면서도 효율적으로 경로 존재 여부를 확인하고 작업을 처리해야 한다.

질의는 두 가지 유형이다:
1. 두 정점을 연결하는 경로가 존재하는지 확인 ((d = 0)).
2. 경로 존재 여부를 확인하고, 존재 여부에 따라 특정 에지를 제거 ((d = 1)).

핵심 아이디어

1. 트리의 구조 표현:

   • 각 정점의 부모를 기반으로 트리를 표현하고, 에지 정보를 유지한다.

2. Heavy-Light Decomposition (HLD):

   • 트리를 체인으로 분리하여 경로 질의를 효율적으로 수행한다.
   • HLD로 트리를 분리한 후, 각 체인의 연결 여부를 세그먼트 트리로 관리한다.

3. 세그먼트 트리:

   • 각 체인에서 에지의 활성 상태를 관리하며 경로 연결 여부를 확인한다.
   • 세그먼트 트리를 사용해 특정 범위의 모든 에지가 활성 상태인지 효율적으로 확인.

4. 질의 처리:

   • 두 정점의 공통 조상을 기준으로 경로를 나누어 질의한다.
   • 경로 존재 여부에 따라 에지를 제거하거나 답을 출력.

코드

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

public class Main {
    static final int MAX = 200001;

    static int N, Q;
    static List<Integer>[] adj;
    static int[] subtreeSize, HLDNum, HLDHead, HLDDepth, HLDParent;
    static int[] tree;
    static int HLDCount = 1;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        PrintWriter pw = new PrintWriter(new OutputStreamWriter(System.out));

        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());

        adj = new ArrayList[MAX];
        for (int i = 0; i < MAX; i++) adj[i] = new ArrayList<>();

        subtreeSize = new int[MAX];
        HLDNum = new int[MAX];
        HLDHead = new int[MAX];
        HLDDepth = new int[MAX];
        HLDParent = new int[MAX];

        int treeHeight = (int) Math.ceil(Math.log(N) / Math.log(2));
        int treeSize = (1 << (treeHeight + 1));
        tree = new int[treeSize];
        Arrays.fill(tree, 1);

        for (int i = 1; i < N; i++) {
            int p = Integer.parseInt(br.readLine());
            adj[p].add(i + 1);
        }

        HLDHead[1] = 1;
        calcSubtreeSize(1);
        heavyLightDecomposition(1, 1);

        while (Q-- > 0) {
            st = new StringTokenizer(br.readLine());
            int b = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());
            int d = Integer.parseInt(st.nextToken());

            boolean result = queryPath(b, c);

            pw.println(result ? "YES" : "NO");

            if (d == 1) {
                if (result) updateTree(1, 1, N, HLDNum[b], 0);
                else updateTree(1, 1, N, HLDNum[c], 0);
            }
        }

        pw.flush();
    }

    static void calcSubtreeSize(int node) {
        subtreeSize[node] = 1;
        for (int child : adj[node]) {
            calcSubtreeSize(child);
            subtreeSize[node] += subtreeSize[child];
            if (subtreeSize[child] > subtreeSize[adj[node].get(0)]) {
                Collections.swap(adj[node], 0, adj[node].indexOf(child));
            }
        }
    }

    static void heavyLightDecomposition(int node, int depth) {
        HLDNum[node] = HLDCount++;
        HLDDepth[node] = depth;

        for (int child : adj[node]) {
            if (child == adj[node].get(0)) {
                HLDHead[child] = HLDHead[node];
                heavyLightDecomposition(child, depth);
            } else {
                HLDHead[child] = child;
                heavyLightDecomposition(child, depth + 1);
            }
        }
    }

    static boolean queryPath(int u, int v) {
        while (HLDHead[u] != HLDHead[v]) {
            if (HLDDepth[HLDHead[u]] < HLDDepth[HLDHead[v]]) {
                int temp = u;
                u = v;
                v = temp;
            }
            if (queryTree(1, 1, N, HLDNum[HLDHead[u]], HLDNum[u]) == 0) return false;
            u = HLDParent[HLDHead[u]];
        }
        if (HLDNum[u] > HLDNum[v]) {
            int temp = u;
            u = v;
            v = temp;
        }
        return queryTree(1, 1, N, HLDNum[u] + 1, HLDNum[v]) == 1;
    }

    static void updateTree(int node, int start, int end, int index, int value) {
        if (index < start || index > end) return;
        if (start == end) {
            tree[node] = value;
            return;
        }

        int mid = (start + end) / 2;
        updateTree(node * 2, start, mid, index, value);
        updateTree(node * 2 + 1, mid + 1, end, index, value);
        tree[node] = tree[node * 2] & tree[node * 2 + 1];
    }

    static int queryTree(int node, int start, int end, int left, int right) {
        if (left > end || right < start) return 1;
        if (left <= start && end <= right) return tree[node];

        int mid = (start + end) / 2;
        return queryTree(node * 2, start, mid, left, right) &
               queryTree(node * 2 + 1, mid + 1, end, left, right);
    }
}

코드 설명

1. DFS를 통한 트리 분할:

• 서브트리 크기를 계산하고, HLD를 통해 체인을 구성한다.

2. HLD를 활용한 경로 질의:

• 두 정점의 경로를 나누어, 체인 단위로 경로의 연결 여부를 확인한다.

3. 세그먼트 트리:

• 체인의 에지 활성 상태를 관리하며, 경로 연결 여부를 효율적으로 확인한다.

4. 질의 처리:

• 질의 결과를 출력하고 필요 시 에지를 제거하여 트리를 업데이트한다.

So...

이 문제는 트리의 구조와 경로 질의를 효율적으로 처리하기 위해 HLD와 세그먼트 트리를 결합한 풀이를 요구한다. 질의 수와 트리 크기가 매우 크기 때문에, 최적화를 통해 효율적으로 문제를 해결하였다.