백준 16121번 사무실 이전

문제 요약

크기 N인 트리가 주어진다. 그리고 모든 간선의 길이는 1이다.

거주지인 정점이 M개, 후보지인 정점이 K개 주어진다.

원하는 것은 모든 거주지와 모든 후보지 간의 거리의 제곱의 합을 구해야 한다.

N은 최대 30만이고 M, K는 둘 다 N만큼 커질 수 있다.

풀이

아무 정점이나 잡아서 그 정점을 루트로 삼은 rooted tree를 생각해보자.

루트의 서브 트리에 속해 있는 거주지와 root와의 거리를 $A_i$라고 하고, 서브트리에 속해 있는 후보지와 root와의 거리를 $B_i$라고 하자.

그러면 거주지와 후보지 간의 경로들 중에서 root를 지나는 경로들의 제곱의 합은 $\sum \sum (A_i+B_j)^2$가 된다.

이를 일일이 다 계산하면 당연히 시간초과를 받는다. 우리가 구해야되는 식을 한번 전개해보자.

$A_i$, $B_j$의 갯수를 $cnt_A, cnt_B$라고 해보자. 그리고 $A_i$의 합은 $sum_A$, $B_j$의 합은 $sum_B$, $A_i^2$의 합은 $sq_A$, $B_j^2$의 합은 $sq_B$라고 하자.
$$
\begin{aligned}
\sum_i \sum_j (A_i^2 + 2A_iB_j + B_j^2) &= \sum_i(cnt_bA_i^2+2A_isum_B+sq_B) \\
&=sq_Acnt_B+2sum_Asum_B+cnt_Asq_B
\end{aligned}
$$
우리가 정의한 $cnt, sum, sq$의 값을 잘 관리해주면 답을 구할 수 있다.

서브트리를 순회하면서 루트와의 거리가 d인 거주지를 만났다면 이제까지 $cnt, sum, sq$를 잘 관리했다면 답에는 $sq_B + cnt_Bd^2+2sum_Bd$를 더해주면 된다. 반대로 루트와의 거리가 d인 후보지를 만났다면 답에 $sq_A + cnt_Ad^2+2sum_Ad$를 답에 더해주면 된다.

답을 갱신하는 과정을 거친 다음엔 $cnt, sq, sum$를 갱신하는 과정을 수행한다. 루트와의 거리가 d인 거주지를 만났으면 $cnt_A$는 1 늘리고, $sum_A$는 d 늘리고, $sq_A$는 $d^2$만큼 늘려주면 된다. 후보지에 대해서도 동일하다.

이제 이 과정을 Centroid Decomposition을 통해서 수행하면 $O(NlogN)$에 문제를 해결할 수 있다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const int MXN = 3e5;
const ll mod = 998244353;
int N, M, K;
int sub_sz[MXN+5];
bool vis[MXN+5], live[MXN+5], cand[MXN+5];
vector<vector<int>> G(MXN+5);
ll ans;
ll cand_cnt, cand_sum, cand_sq;
ll live_cnt, live_sum, live_sq;

int get_size(int cur, int par) {
    sub_sz[cur] = 1;
    for(int &nxt:G[cur]) {
        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(int &nxt:G[cur]) {
        if(nxt == par || vis[nxt]) continue;
        if(sub_sz[nxt] > thr) return get_cent(nxt, cur, thr);
    }
    return cur;
}

void query(int cur, int par, ll depth) {
    ll tmp = 0;
    if(cand[cur]) {
        tmp += live_sq;
        tmp %= mod;
        tmp += (live_cnt*depth % mod) *depth % mod;
        tmp %= mod;
        tmp += (2LL*live_sum%mod)*depth % mod;
        tmp %= mod;
    }
    if(live[cur]) {
        tmp += cand_sq;
        tmp %= mod;
        tmp += (cand_cnt*depth % mod) *depth % mod;
        tmp %= mod;
        tmp += (2LL*cand_sum%mod)*depth%mod;
        tmp %= mod;
    }
    ans += tmp; ans %= mod;
    for(int &nxt:G[cur]) {
        if(nxt == par || vis[nxt]) continue;
        query(nxt, cur, depth+1);
    }
}

void update(int cur, int par, ll depth) {
    if(cand[cur]) {
        cand_cnt++;
        cand_sq += depth*depth%mod;
        cand_sq %= mod;
        cand_sum += depth;
        cand_sum %= mod;
    }
    if(live[cur]) {
        live_cnt++;
        live_sq += depth*depth%mod;
        live_sq %= mod;
        live_sum += depth;
        live_sum %= mod;
    }
    for(int &nxt:G[cur]) {
        if(nxt == par || vis[nxt]) continue;
        update(nxt, cur, depth+1);
    }
}

void init() {
    cand_cnt = cand_sq = cand_sum = 0;
    live_cnt = live_sq = live_sum = 0;
}

void dnc(int cur) {
    int thr = get_size(cur, 0);
    int cent = get_cent(cur, 0, thr/2);
    init();
    vis[cent] = true;
    if(cand[cent]) ++cand_cnt;
    if(live[cent]) ++live_cnt;
    for(int &nxt:G[cent]) {
        if(!vis[nxt]) {
            query(nxt, cent, 1);
            update(nxt, cent, 1);
        }
    }
    for(int &nxt:G[cent]) {
        if(!vis[nxt]) dnc(nxt);
    }
}

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);
    }
    cin >> M;
    for(int i=0;i<M;++i) {
        int x; cin >> x;
        live[x] = true;
    }
    cin >> K;
    for(int i=0;i<K;++i) {
        int x; cin >> x;
        cand[x] = true;
    }
    dnc(1);
    cout << ans << '\n';
    return 0;
}