[백준 13512] 트리와 쿼리 3

문제

www.acmicpc.net/problem/13512

 

알고리즘

HLD

 

풀이

트리 위에서 두 가지 쿼리를 진행해야 합니다. 첫 번째는 $i$번째 정점의 색을 바꾸는 것. 두 번째는 1번 정점에서부터 $v$정점까지 가는 경로 중 첫 번째로 만나는 검은 정점의 번호를 출력하는 것입니다.

 

트리에서는 문제를 풀기가 어려우므로 HLD를 이용해서 쭉 펴도록 합시다. 일자 상태로 만들었다면 세그먼트 트리의 사용이 가능해집니다. 세그먼트 트리에서 인덱스가 작을수록 트리 위에서는 루트에 가깝습니다. 이러한 이유 때문에 $query$함수에서 세그먼트 트리 쿼리를 처리할 때, 왼쪽에 검은 정점이 있다면 바로 리턴을 했습니다. 오른쪽에 있는 구간은 왼쪽보다 루트로부터 멀기 때문입니다(처음으로 만나는 검은 정점이 아님).

 

세그먼트 트리의 리프노드에는 흰 정점이라면 0을 검은 정점이라면 해당 번호를 적어놓고 풀었습니다.

 

코드

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; ++i)
#define REP(i, n) for (int i = 1; i <= n; ++i)
using namespace std;

#define left (i << 1)
#define right (i << 1 | 1)

int N, M, cnt;
vector<vector<int>> adj;
vector<int> p, idx, st, top, tree;

void dfs(int here) {
    st[here] = 1;
    for (auto& there : adj[here]) {
        if (there ^ p[here]) {
            p[there] = here;
            dfs(there);
            st[here] += st[there];
            int& HE = adj[here][0];
            if (HE == p[here] || st[there] > st[HE])
                swap(HE, there);
        }
    }
}

void hld(int here) {
    idx[here] = cnt++;
    for (auto there : adj[here]) {
        if (there ^ p[here]) {
            top[there] = there == adj[here][0] ? top[here] : there;
            hld(there);
        }
    }
}

void update(int i) {
    tree[idx[i] + N] = tree[idx[i] + N] ? 0 : i + 1;
    i = idx[i] + N;
    for (i >>= 1; i; i >>= 1) {
        tree[i] = tree[left] ? tree[left] : (tree[right] ? tree[right] : 0);
    }
}

int query(int l, int r) {
    int ret = 0;
    l += N;
    r += N;
    for (; l <= r; l >>= 1, r >>= 1) {
        if ((l & 1) && tree[l++]) {
            return tree[l - 1];
        }
        if (!(r & 1) && tree[r--]) {
            ret = tree[r + 1];
        }
    }
    return ret;
}

int solve(int u, int v) {
    int ret = -1;
    while (top[u] ^ top[v]) {
        if (st[top[u]] < st[top[v]])
            swap(u, v);
        int cand = query(idx[top[v]], idx[v]);
        ret = cand ? cand : ret;
        v = p[top[v]];
    }
    if (idx[u] < idx[v])
        swap(u, v);
    int cand = query(idx[v], idx[u]);
    ret = cand ? cand : ret;
    return ret;
}

int main() {
#ifndef ONLINE_JUDGE
    freopen("in", "r", stdin);
    freopen("out", "w", stdout);
#endif
    cin.tie(NULL);
    cout.tie(NULL);
    ios::sync_with_stdio(false);

    cin >> N;
    adj.resize(N);
    p = idx = top = st = vector<int>(N);
    rep(i, N - 1) {
        int u, v;
        cin >> u >> v;
        --u, --v;
        adj[u].emplace_back(v);
        adj[v].emplace_back(u);
    }
    dfs(0);
    hld(0);
    tree.resize(N << 1);
    cin >> M;
    while (M--) {
        int a, b;
        cin >> a >> b;
        --b;
        if (a == 1) {
            update(b);
        } else {
            cout << solve(0, b) << '\n';
        }
    }
    return 0;
}