[C++] 백준 13309: 트리

1. 문제 출처

 

https://www.acmicpc.net/problem/13309

13309번: 트리

 

 

 

 

 

 

2. 문제 접근

 

특정 구간내에 그래프가 연결되어 있는지의 정보를 O(log N)의 시간만에 알아내기 위해 세그먼트 트리 자료구조를 사용해야 했다. 하지만, 트리 구조이기 때문에 트리구조를 체인으로 분할 하는 heavy-light-decomposition 을 적용하여 문제를 해결해야 한다.

 

 

 

 

먼저 트리의 루트는 항상 1이라고 고정된 상태이기 때문에, 1을 루트로 하여 각 노드의 서브트리 크기를 구한다. 이후에 heavy light decomposition 과정에서 서브트리의 크기가 가장 큰 노드가 연결 리스트의 가장 앞에 오도록 하여 heavy subtree와 light subtree를 분할하였다.

 

 

 

 

질의를 처리하는 과정에서는 세그먼트 트리를 사용했는데, 해당 구간에 하나라도 끊긴 부분이 있어서는 안되기 때문에 AND 연산자를 사용한 세그먼트 트리를 구성하였다. 쿼리를 진행하는 과정에서는 최소 공통 조상을 찾아서 올라가는 알고리즘을 추가하여 정확한 구간을 찾을 수 있도록 만들었다.

 

 

 

 

 

 

3. 코드

#include <bits/stdc++.h>
#define ll long long
#define pii pair<int, int>
#define MAX 200001

using namespace std;

int N;  // number of nodes in tree
int Q;  // number of queries

vector<int> tree;           // AND operation segment tree
vector<int> adj[MAX];       // adjacency list

int subtree_size[MAX];      // subtree size
int HLD_num[MAX];           // heavy-light-decomposition id
int HLD_head[MAX];          // top of group
int HLD_depth[MAX];         // end of group
int HLD_parent[MAX];        // parent group
int HLD_cnt = 1;            // heavy-light-decomposition id numbering


// update segment tree
void update_tree(int node, int start, int end, int index, int value) {
    // out of range
    if (index < start || end < index) return;

    tree[node] = value;

    // not leaf node
    if (start != end) {
        int mid = (start + end) / 2;
        update_tree(node * 2, start, mid, index, value);
        update_tree(node * 2 + 1, mid + 1, end, index, value);
        tree[node] = tree[node * 2] & tree[node * 2 + 1];
    }
}


// query segment tree
int query_tree(int node, int start, int end, int left, int right) {
    // out of range
    if (right < start || end < left) return 1;
    // query range in node range
    if (left <= start && end <= right) return tree[node];

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


void calc_subtree_size(int here, int level) {
    subtree_size[here] = 1;     // self

    for (int &there : adj[here]) {
        calc_subtree_size(there, level + 1);
        subtree_size[here] += subtree_size[there];

        // for heavy light decomposition, place heaviest subtree to front.
        if (subtree_size[there] > subtree_size[adj[here][0]] || adj[here][0] == there) {
            swap(adj[here][0], there);
        }
    }
}


void heavy_light_decomposition(int here, int depth) {
    HLD_num[here] = HLD_cnt++;
    HLD_depth[here] = depth;

    for (int there : adj[here]) {
        // heaviest subtree
        if (there == adj[here][0]) {
            HLD_head[there] = HLD_head[here];
            HLD_parent[there] = HLD_parent[here];
            heavy_light_decomposition(there, depth);
        }

        else {
            HLD_head[there] = there;
            HLD_parent[there] = here;
            heavy_light_decomposition(there, depth + 1);
        }
        
        // update segment tree with AND(&) calculation
        update_tree(1, 1, N, HLD_num[there], 1);
    }
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    
    cin >> N >> Q;

    int tree_height = (int)ceil(log2(N));
    int tree_size = (1 << (tree_height + 1));
    tree.resize(tree_size, 1);

    for (int i = 1; i < N; ++i) {
        int p; cin >> p;
        adj[p].push_back(i + 1);
    }

    HLD_head[1] = 1;
    calc_subtree_size(1, 1);
    heavy_light_decomposition(1, 1);

    while (Q--) {
        bool result = true;
        int b, c, d; cin >> b >> c >> d;
        int B = b, C = c;

        if (HLD_depth[b] > HLD_depth[c]) swap(b, c);

        while (HLD_depth[b] < HLD_depth[c]) {
            result &= query_tree(1, 1, N, HLD_num[HLD_head[c]], HLD_num[c]);
            c = HLD_parent[c];
        }

        while (HLD_head[b] != HLD_head[c]) {
            result &= query_tree(1, 1, N, HLD_num[HLD_head[b]], HLD_num[b]);
            result &= query_tree(1, 1, N, HLD_num[HLD_head[c]], HLD_num[c]);

            b = HLD_parent[b];
            c = HLD_parent[c];
        }

        if (HLD_num[b] > HLD_num[c]) swap(b, c);
        result &= query_tree(1, 1, N, HLD_num[b] + 1, HLD_num[c]);

        cout << (result ? "YES" : "NO") << "\n";

        if (d) {
            if (result) update_tree(1, 1, N, HLD_num[B], 0);
            else update_tree(1, 1, N, HLD_num[C], 0);
        }
    }
    
    return 0;
}