백준 13513번 트리와 쿼리 4

문제 요약

크기 N인 트리가 주어진다. 처음 정점의 색은 전부 흰색이다. 다음과 같은 쿼리를 처리하자.

1. i번 정점의 색을 바꾼다. 흰색이면 검은색으로, 검은색이면 흰색으로
2. 모든 흰색 정점 쌍 중에서 그 정점간의 거리가 가장 긴 거리를 출력하자.

풀이

내가 짰던 코드를 보기만 해도 머리가 아파온다.
가능하면 백준 13431번 트리 문제을 풀고 이 문제를 푸는 것을 추천한다.

이 문제가 훨씬 간단한거 같다.

 

일단 전체에서 가장 긴 흰색 정점간의 거리를 처리해야 한다.

위 사진에서 선언한 컨테이너들의 역할부터 차례차례 소개하겠다.

 

subtree[x]는 센트로이드 트리를 만들었을 때, x의 서브트리들에 속해 있는 흰색 정점과 x와의 거리를 각 서브트리별로 들고 있다.

 

sub_dm[x]는 x의 서브트리에 속해 있는 흰색 정점들 중 x와의 거리가 가장 먼 흰색 정점과의 거리를 들고 있다. 이 때, 관리하는 거리는 x의 서브트리 당 하나이다. 즉, sub_dm[x]는 x의 서브트리 갯수만큼 거리를 들고 있다.

sub_dm[x]에서 가장 큰 두 값을 합쳐주면 x를 루트로 하는 서브트리에서 x를 지나는 흰색 정점간의 경로 중 가장 긴 경로의 길이가 된다.

 

ans는 sub_dm[x]에서 경로가 생기면 ans에 저장하는 용도이다. 2번 쿼리가 들어오면 ans에서 가장 큰 값을 출력하면 된다.

 

dp는 각 노드 별로 현재 sub_dm[x]에서 생긴 경로들을 확인하는 용도이다. 지금 보니 필요 있나 싶기도 하다.

 

어쨌든 2번 쿼리는 ans에서 가장 큰값을 출력하면 끝난다.

 

1번 쿼리가 들어왔을 때 업데이트를 어떻게 하는지 알아보자.

먼저 i번 노드부터 부모를 타고 쭉 올라간다. 타고 올라가면서 i번 노드가 속한 서브트리를 prev라고 하자.

부모 노드 p와 i번 노드 간의 거리를 구하고 이 거리를 subtree[p][prev]에 추가하거나 삭제해야 한다.

 

흰색 정점으로 바뀌는 경우라서 추가해야 할 경우에 이 거리가 만약 subtree[p][prev]에서 가장 큰 값이라면 sub_dm[p]도 갱신을 해줘야 한다. 그리고 sub_dm[p]가 갱신되었으니 경로가 존재하는지 확인하고 ans도 갱신한다.

 

검은색 정점으로 바뀌는 경우라서 삭제해야 할 경우 이 거리가 subtree[p][prev]에서 제일 큰 값이었다면 sub_dm[p]가 바뀌게 되고 ans도 갱신해준다.

 

이거로 시간복잡도 $O((N+Q)log^2N)$에 문제를 해결할 수 있다.

#include<bits/stdc++.h>
#define MAXN 100100
#define LOGN 20
using namespace std;
using ll = long long;
const int INF = -1e9;
int N, M, timer;
int euler[2*MAXN], lev[MAXN];
ll h[MAXN];
vector<vector<int>> idx(MAXN*2);
vector<vector<pair<int,ll>>> G(MAXN);

void dfs(int cur, int par, int l, ll d){
    lev[cur] = l;
    h[cur] = d;
    euler[++timer] = cur;
    idx[cur].push_back(timer);
    for(const auto &p : G[cur]) {
        int v = p.first;
        ll w = p.second;
        if(v == par) continue;
        dfs(v, cur, l+1, d+w);
        euler[++timer] = cur;
        idx[cur].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] = {lev[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;
}

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

vector<int> sub_sz, parent;
vector<bool> vis, color;

int get_size(int cur, int par) {
    sub_sz[cur] = 1;
    for(auto &p:G[cur]) {
        int nxt = p.first;
        if(nxt == par || vis[nxt]) continue;
        sub_sz[cur] += get_size(nxt, cur);
    }
    return sub_sz[cur];
}

int get_cent(int cur, int par, int thr) {
    for(auto &p:G[cur]) {
        int nxt = p.first;
        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_size(cur, -1);
    int cent = get_cent(cur, -1, thr/2);
    vis[cent] = true;
    if(par == -1) parent[cent] = 0;
    else parent[cent] = par;
    for(auto &p:G[cent]) {
        int nxt = p.first;
        if(!vis[nxt]) make_tree(nxt, cent);
    }
}

vector<unordered_map<int,multiset<int>>> subtree;
vector<multiset<int>> sub_dm;
multiset<int> ans;
vector<int> dp;
int cnt;

void update_ans(int idx) {
    if(sub_dm[idx].size() < 2) {
        if(dp[idx] != INF) ans.erase(ans.find(dp[idx]));
        dp[idx] = INF;
        return ;
    }
    int tmp = *(sub_dm[idx].rbegin()) + *next(sub_dm[idx].rbegin());
    if(dp[idx] == INF) {
        ans.insert(tmp);
        dp[idx] = tmp;
    }
    else if(dp[idx] != tmp) {
        ans.erase(ans.find(dp[idx]));
        ans.insert(tmp); dp[idx] = tmp;
    }
}

void update(int cur) {
    if(color[cur]) --cnt;
    else ++cnt;
    color[cur] = !color[cur];
    int prev = cur;
    int p = cur;
    while(p) {
        int d = 0;
        if(p != cur) d = dist(p, cur);
        auto &ms = subtree[p][prev];
        if(color[cur]) {
            if(ms.lower_bound(d) == ms.end()) {
                if(!ms.empty()) sub_dm[p].erase(sub_dm[p].find(*(ms.rbegin())));
                ms.insert(d); sub_dm[p].insert(d);
                update_ans(p);
            }
            else {
                ms.insert(d);
            }
        }
        else {
            int mx = *(ms.rbegin());
            ms.erase(ms.find(d));
            if(ms.lower_bound(mx) == ms.end()) {
                sub_dm[p].erase(sub_dm[p].find(mx));
                if(!ms.empty()) sub_dm[p].insert(*(ms.rbegin()));
                update_ans(p);
            }
        }
        prev = p;
        p = parent[p];
    }
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N;
    for(int i=1;i<N;++i) {
        int u, v, w; cin >> u >> v >> w;
        G[u].emplace_back(v,w);
        G[v].emplace_back(u,w);
    }
    dfs(1, -1, 0, 0);
    sparsetable_build();
    vis.resize(N+1); color.resize(N+1); sub_sz.resize(N+1); parent.resize(N+1);
    subtree.resize(N+1); dp.resize(N+1, INF); sub_dm.resize(N+1);
    make_tree(1, -1);
    for(int i=1;i<=N;++i) {
        subtree[i][i] = multiset<int>();
        if(parent[i]) subtree[parent[i]][i] = multiset<int>();
    }
    for(int i=1;i<=N;++i) update(i);
    cin >> M;
    for(int i=0;i<M;++i) {
        int q, x;
        cin >> q;
        if(q == 1) {
            cin >> x; update(x);
        }
        else {
            if(!cnt) cout << -1 << '\n';
            else if(cnt == 1) cout << 0 << '\n';
            else {
                int tmp = *(ans.rbegin());
                if(tmp < 0) tmp = 0;
                cout << tmp << '\n';
            }
        }
    }
    return 0;
}