백준 13431번 트리 문제

문제 요약

크기가 N인 트리가 주어진다. 처음에 모든 정점은 흰색이다. 다음과 같은 쿼리를 처리해야 한다.

1. x번 정점을 파란색으로 색칠한다.
2. x번 정점과 모든 파란 정점과의 거리의 합을 계산한다.

풀이

다행히도 파란색이 흰색으로 바뀌는 쿼리가 없다.

각 정점에 대해서 다음과 같은 dp값들을 정의하자.
$$
dp(i) = \text{i번 노드를 루트로 하는 서브트리에서 i번 노드부터 파란 정점까지의 거리의 합}
$$
이제 1번 쿼리가 들어오면 x번 정점의 부모를 쭉 타고 올라가면서 부모의 dp값에 부모와 x번 정점간의 거리를 더해주면 된다.

그러나, 부모를 타고 가는 횟수는 최대 N-1번이고 트리가 선형으로 존재하고 리프노드부터 순서대로 타고 올라가면 $O(N^2)$으로 시간 초과가 난다.

이를 막기 위해서 센트로이드 트리 상에서 위와 같은 트리 dp를 진행한다고 하자.

이제 2번 쿼리를 처리해야 한다. 2번 쿼리도 x번 정점에서부터 부모로 올라가면서 거리의 합을 구하고 싶지만 부모의 dp값에는 x번 정점이 속해있는 서브트리에 대한 값들도 포함이 되어 있다. 이를 잘 처리해줘야 한다.

즉, $dp(par)$ 값에서 x가 속해있는 서브트리에 있는 파란 정점들과 par과의 거리를 $dp(par)$에서 빼준 값을 쿼리의 답에 더해야 한다.

이를 위해서 다음과 같은 두가지 값을 더 관리해줘야 한다.
$$
cnt(i) = \text{i번 노드를 루트로 하는 서브트리에 속하는 파란 정점의 수}
$$

$$
dp2(p,x)=\text{정점 p에서 p의 서브트리 x에 속하는 파란 정점들까지의 거리의 합}
$$

이런 값들이 잘 관리가 되어있다면 $x$가 쿼리로 들어왔을 때 부모로 타고 올라가면서 값을 적절히 빼줘서 답을 구하면 된다.

$prev$가 $x$가 속해있는 $p$의 서브트리라고 하면 부모를 타고 올라가면서 $dp(p) - dp2(p,prev)+(cnt(p)-cnt(prev))dist(p,x)$를 더해나가면 된다.

이것으로 문제를 $O(Qlog^2N+NlogN)$에 해결이 가능하다.

여담으로 이 문제는 백준 12023번 문제와 동일한 문제이다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const int MAXN = 1e5+5;
const int LOGN = 20;
vector<int> parent;
int sub_sz[MAXN];
bool vis[MAXN], color[MAXN];
ll lev[MAXN], h[MAXN];
vector<vector<pair<int,ll>>> G(MAXN);
int N, Q, timer;
int euler[2*MAXN];
vector<vector<int>> idx(MAXN*2);
vector<pair<int,ll>> dp;
unordered_map<int, ll> mp[MAXN];

void dfs(int cur, int par, int l, ll d) {
    h[cur] = d;
    lev[cur] = l;
    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;
}

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

int get_cent(int cur, int par, int thr) {
    for(const auto &p:G[cur]) {
        int v = p.first;
        if(v == par || vis[v]) continue;
        if(sub_sz[v] > thr) return get_cent(v, 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(const auto &p:G[cent]) {
        int v = p.first;
        if(vis[v]) continue;
        make_tree(v, cent);
    }
}

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

int cnt;

void update(int cur) {
    if(color[cur]) return ;
    ++cnt;
    color[cur] = true;
    dp[cur].first++;
    int p = parent[cur];
    int prev = cur;
    while(p) {
        ll d = dist(cur, p);
        dp[p].second += d;
        mp[p][prev] += d;
        dp[p].first++;
        prev = p;
        p = parent[p];
    }
}

ll query(int cur) {
    ll ret = dp[cur].second;
    int p = parent[cur];
    int prev = cur;
    while(p) {
        int cnt = dp[p].first - dp[prev].first;
        ret += (dp[p].second - mp[p][prev]);
        ret += cnt * dist(p, cur);
        prev = p;
        p = parent[p];
    }
    return ret;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N >> Q;
    parent.resize(N+1); dp.resize(N+1);
    for(int i=1;i<N;++i) {
        int u, v, w;
        cin >> u >> v >> w;
        ++u; ++v;
        G[u].push_back({v,w});
        G[v].push_back({u,w});
    }
    dfs(1, -1, 0, 0);
    sparsetable_build();
    make_tree(1, -1);
    for(int i=0;i<Q;++i) {
        int q, x;
        cin >> q >> x;
        ++x;
        if(q == 1) update(x);
        else cout << query(x) << '\n';
    }
    int root;
    for(int i=1;i<=N;++i) if(!parent[i]) root = i;
    assert(dp[root].first == cnt);
    return 0;
}