백준 5820번 경주

문제 요약

크기가 N인 트리가 주어진다. 이 떄, 트리에 존재하는 경로 중에서 길이가 K면서 이루는 간선의 수가 가장 작은 경로의 간선 수를 출력하라. 만약에 길이가 K인 경로가 없다면 -1을 출력한다.

N은 최대 10만, K는 최대 100만이다.

풀이

rooted tree에서 해당 트리에서 root와의 거리가 d인 정점 중에서 그 경로의 간선 수가 가장 적은 경우의 간선 수를 $min\_depth(d)$라고 합시다.

그러면 이러한 $min\_depth$는 각 서브트리를 순회하면서 갱신이 가능하고, 서브트리를 순회하면서 root와의 거리가 d인 정점을 만났을 때, $min\_depth(K-d)$가 초기값이 아니라 갱신이 되어 있는 상태라면 서로 다른 서브트리에서 root를 지나면서 길이가 K인 경로를 발견한 것입니다.

이 때, 답을 갱신해주도록 합니다.

이런 식으로 답을 갱신하는 과정을 거친 다음에는 $min\_depth$를 갱신하는 과정을 진행합니다.

이를 모든 정점에 대해서 일일이 하면 $O(N^2)$으로 시간초과를 받지만 Centroid Decomposition을 사용하면 $O(NlogN)$에 해결이 가능합니다. $min\_depth$를 초기화 할 땐 갱신됐던 값들만 초기화하도록 주의합시다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const int MXN = 200000;
const int MXK = 1000000;
const int INF = 1e9;
int sub_sz[MXN+5];
bool vis[MXN+5];
int min_depth[MXK+5];
vector<int> cur_tree;
vector<vector<pair<int,int>>> G(MXN+5);
int N, K;

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

// return centroid of subtree which has cur as root
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;
}

int solve(int cur, int par, int dist, int depth) {
    int ret = INF;
    if(dist > K) return ret;
    ret = min(ret, min_depth[K-dist]+depth);
    for(const auto &p:G[cur]) {
        int v = p.first;
        int w = p.second;
        if(v == par || vis[v]) continue;
        ret = min(ret, solve(v, cur, dist+w, depth+1));
    }
    return ret;
}

void update(int cur, int par, int dist, int depth) {
    if(dist > K) return ;
    min_depth[dist] = min(min_depth[dist], depth);
    cur_tree.push_back(dist);
    for(const auto &p:G[cur]) {
        int v = p.first;
        int w = p.second;
        if(v == par || vis[v]) continue;
        update(v, cur, dist+w, depth+1);
    }
}

int dnc(int cur) {
    int thr = get_size(cur, -1);
    int ct = get_cent(cur, -1, thr/2); // Centroid Found
    vis[ct] = 1;                    // remove centroid from tree
    int ret = INF;
    for(const auto &a:cur_tree) min_depth[a] = INF;
    cur_tree.clear();
    min_depth[0] = 0;
    for(const auto &p:G[ct]) {      // now we can divide tree into subtrees <= thr
        int v = p.first;
        int w = p.second;
        if(!vis[v]) {
            ret = min(ret,solve(v, ct, w, 1));      // conquer subtrees
            update(v, ct, w, 1);
        }
    }
    for(const auto &p:G[ct]) {
        int v = p.first;
        if(vis[v]) continue;
        ret = min(ret, dnc(v));
    }
    return ret;
}

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