백준 13514번 트리와 쿼리 5

문제 요약

크기 N의 트리가 주어진다. 각 정점은 1번부터 N번으로 숫자가 매겨져 있고, 처음의 모든 정점의 색은 검은색이다. 모든 간선의 길이는 1이며, 무방향 간서이다.

다음과 같은 쿼리를 처리 해야한다.

1. i번 정점의 색을 바꾼다. 흰색이면 검은색으로, 검은색이면 흰색으로
2. i번과 가장 가까운 흰색 정점과의 거리를 출력한다.

N은 최대 10만이고, 쿼리 수도 최대 10만이다.

풀이

트리 DP 형식으로 생각을 해봅시다.
$$
dp(i)=\text{i번을 루트로 하는 서브트리에서 i번과 가장 가까운 흰색 정점의 거리}
$$
이런 값들을 잘 관리할 수 있다면 2번 쿼리가 들어왔을 때, 쿼리가 들어온 정점의 부모를 타고 가면서 (부모와 정점 간의 거리) + 부모의 dp값을 비교하면서 답을 찾을 수 있을 겁니다

그러나, 부모를 타고 가는 횟수가 최대 N-1번 있을 수 있기 때문에 dp를 기막히게 관리했다 쳐도 시간초과가 납니다.

이럴 때 사용해볼 만한 것이 센트로이드 트리입니다. 센트로이드 트리는 높이가 $O(logN)$이기 때문에 최대 $O(logN)$개 만큼의 부모를 살펴보게 될 것이고 부모와 쿼리로 들어온 정점 간의 거리를 구하는 데에 $O(logN)$만큼 걸리니 쿼리당 $O(log^2N)$에 해결이 가능할 거 같습니다.

그러면 $dp$값을 관리하는 방법을 생각해봐야 합니다. 흰색에서 검은색으로, 검은색에서 흰색으로도 갈 수 있기 때문에 $i$번을 루트로 하는 서브트리의 점들이 $dp(i)$에 고려되었다가 빠졌다가 하는게 되야 합니다.

이는 각 정점마다 pq나 multiset을 만들고 1번 쿼리로 어떤 정점의 색이 바뀌었다면 해당 정점부터 센트로이드 트리에서 부모를 타고 올라가면서 부모의 pq나 multiset를 수정해주면 됩니다.

정점 $v$가 흰색에서 검은색으로 간다면 $v$부터 시작해서 부모를 타고 올라가면서 $dp$값을 수정해줍니다.

이 과정에서 수정하게 되는 부모의 개수는 $O(logN)$이고, $dp$값을 수정하는 데에 $O(logN)$, 그리고 여기서도 쿼리로 들어온 정점과 부모간의 거리를 구하는 데에 $O(logN)$이 걸리므로 쿼리를 처리하는 데에 $O(log^2N)$이 걸립니다.

이걸로 총 시간복잡도는 $O((N+Q)log^2N)$이 됩니다.

#include<bits/stdc++.h>
using namespace std;
const int INF = 1e9;
const int MAXN = 1e5+5;
const int LOGN = 17;
bool color[MAXN], vis[MAXN];
int parent[MAXN], sub_sz[MAXN];
int N, M, timer;
int depth[MAXN];
priority_queue<pair<int,int>, vector<pair<int,int>>, greater<pair<int,int>>> dp[MAXN];
int euler[2*MAXN];
vector<vector<int>> G(MAXN), idx(MAXN*2);

void dfs(int x, int pa, int l){
    depth[x] = l;
    euler[++timer] = x;
    idx[x].push_back(timer);
    for(int e : G[x]) if( e != pa ){
        dfs(e, x, l+1);
        euler[++timer] = x;
        idx[x].push_back(timer);    
    }
}

int pw2[LOGN], lg2[MAXN*2];
pair<int,int> ST[LOGN][MAXN*2];

void sparsetable_build(){
    // calculate pw2 array
    pw2[0] = 1;
    for(int k=1;k<LOGN;k++) pw2[k] = pw2[k-1] * 2;

    // calculate lg2 array
    memset(lg2, -1, sizeof lg2);
    for(int k=0;k<LOGN;k++) if( pw2[k] < MAXN*2 ) lg2[pw2[k]] = k;
    for(int i=1;i<MAXN*2;i++) if( lg2[i] == -1 ) lg2[i] = lg2[i-1];

    // calculate sparse table
    for(int i=1;i<=2*N-1;i++) ST[0][i] = {depth[euler[i]], euler[i]};

    for(int k=1;k<LOGN;k++){
        for(int i=1;i<=2*N-1;i++){
            if( i + pw2[k-1] > 2*N-1 ) continue;
            ST[k][i] = min(ST[k-1][i], ST[k-1][i+pw2[k-1]]);
        }
    }    
}

int LCA(int u, int v){
    int l = idx[u][0], r = idx[v][0];
    if( l > r ) swap(l,r);
    int k = lg2[r-l+1];
    return min(ST[k][l], ST[k][r-pw2[k]+1]).second;
}

int get_sz(int cur, int par) {
    int &ret = sub_sz[cur];
    ret = 1;
    for(const int &nxt:G[cur]) {
        if(nxt == par || vis[cur]) continue;
        ret += get_sz(nxt, cur);
    }
    return ret;
}

int get_cent(int cur, int par, int thr) {
    for(const int &nxt:G[cur]) {
        if(nxt == par || vis[nxt]) continue;
        if(sub_sz[nxt] > thr) return get_cent(nxt, cur, thr);
    }
    return cur;
}

void make_tree(int cur, int par) {
    int thr = get_sz(cur, -1);
    int cent = get_cent(cur, -1, thr/2);
    if(par == -1) parent[cent] = 0;
    else parent[cent] = par;
    vis[cent] = par;
    for(const int &nxt:G[cent]) if(!vis[nxt]) make_tree(nxt, cent);
}


int dist(int u, int v) {
    return depth[u] + depth[v] - 2*depth[LCA(u,v)];
}

void update(int v) {
    color[v] = !color[v];
    int idx = v;
    while(idx) {
        dp[idx].push({dist(v,idx), v});
        idx = parent[idx];
    }
}

int query(int v) {
    int ret = INF;
    int idx = v;
    while(idx) {
        auto &pq = dp[idx];
        while(!pq.empty()) {
            if(!color[pq.top().second]) pq.pop();
            else break;
        }
        int tmp = INF;
        if(!pq.empty()) tmp = pq.top().first + dist(v, idx);
        ret = min(ret, tmp);
        idx = parent[idx];
    }
    if(ret == INF) ret = -1;
    return ret;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N;
    for(int i=1;i<N;++i) {
        int u, v;
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, -1, 0);
    sparsetable_build();
    make_tree(1, -1);
    cin >> M;
    for(int i=0;i<M;++i) {
        int q, v;
        cin >> q >> v;
        if(q == 1) update(v);
        else cout << query(v) << '\n';
    }
    return 0;
}