백준 20297번 Confuzzle

문제 요약

크기가 N인 트리가 주어진다. 각 정점에는 1부터 N까지의 숫자 중 하나가 적혀있다. 각 간선의 길이는 1이다.

이 때, 같은 번호가 적혀있는 정점 간의 거리 중 제일 짧은 거리를 구해야 한다.

N은 최대 10만이다.

풀이

어떤 정점 하나를 잡았다고 하자. 이 정점을 루트로 보면 여러개의 서브트리가 생긴다.

루트 정점에서 특정 번호까지의 최단 거리를 $min\_depth(i)$라고 하자.

그러면 각 서브트리를 순회하면서 $i$가 적혀있는 정점을 만나면 $min\_depth$를 갱신해주고, 만약에 $min\_depth$가 갱신이 되어 있는 상태라면 답도 같이 갱신합니다.

하나의 정점을 루트로 삼고 진행하면 이 과정은 서로 다른 서브트리에서 루트를 지나는 경로들을 봐준 것이 됩니다.

이제 모든 정점에서 이 과정을 수행하면 되는데 단순하게 수행하면 $O(N^2)$이 걸립니다.

Centroid Decomposition을 통해서 센트로이드를 잡고 해당 트리에서 위 과정을 수행하는 과정을 반복해 나가면 $O(NlogN)$으로 수행 가능합니다.

$min\_depth$를 초기화 할 때, 이전에 갱신이 있었던 부분만 초기화해줘야 시간초과를 피할 수 있습니다.

#include<bits/stdc++.h>
using namespace std;
const int MAXN = 1e5+5;
const int INF = 1e9;
int sub_sz[MAXN];
int min_depth[MAXN];
int a[MAXN];
vector<vector<int>> G(MAXN);
bool vis[MAXN];
vector<int> cur_tree;
int N;

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

int query(int cur, int par, int depth) {
    int ret = INF;
    if(min_depth[a[cur]] != INF) ret = min(ret, min_depth[a[cur]]+depth);
    for(const int &v:G[cur]) {
        if(v == par || vis[v]) continue;
        ret = min(ret, query(v, cur, depth+1));
    }
    return ret;
}

void update(int cur, int par, int depth) {
    min_depth[a[cur]] = min(min_depth[a[cur]], depth);
    cur_tree.push_back(a[cur]);
    for(const int &v:G[cur]) {
        if(v == par || vis[v]) continue ;
        update(v, cur, depth+1);
    }
}

int dnc(int cur) {
    int thr = get_size(cur, 0);
    int ct = get_cent(cur, 0, thr/2);
    for(const int &v:cur_tree) min_depth[v] = INF;
    cur_tree.clear();
    vis[ct] = true;
    min_depth[a[ct]] = 0;
    cur_tree.push_back(a[ct]);
    int ret = INF;
    for(const int &v:G[ct]) {
        if(!vis[v]) {
            ret = min(ret, query(v, ct, 1));
            update(v, ct, 1);
        }
    }
    for(const int &v:G[ct]) {
        if(!vis[v]) ret = min(ret, dnc(v));
    }
    return ret;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N;
    for(int i=1;i<=N;++i) cin >> a[i];
    for(int i=1;i<N;++i) {
        int u, v; cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    for(int i=1;i<=N;++i) min_depth[i] = INF;
    cout << dnc(1) << '\n';
    return 0;
}