yukicoder No.922 東北きりきざむたん

問題ページ
これは を最小化するように空港を設置する問題である．もし と が同じ木の頂点であれば，どのように空港を配置しても2点間の距離は変化しない．木の2点間の距離はLCAなどを用いることによって で求めることができる．

ここからは と が違う木に属する場合についてのみ考える． が属する木の空港の位置を とすると，答えは となる． 番目の木に属するクエリの頂点の集合を とする．答えにこの木が寄与するのは の分であり，空港の位置 を変えることで最小化したい．

もし木上の2点間の距離の和でなく数列上のマンハッタン距離の和を最小化する問題であれば は中央値を選択すればよい．木における中央値のような概念は重心である．これらは 中央値/重心 から値を動かしても答えが改善されないということから証明ができる．ここでいう重心は，「木を分断したときにその部分木のサイズが全て半分以下になるような頂点」という一般的なものそのままとは少し違い，「木を分断したときにその部分木に属するクエリの頂点集合のサイズが全て元の半分以下になるような頂点」を指す．

各木において重心を求め，重心と木に属するクエリの頂点の距離の和を求めればよい．LCA等を用いて2点間の距離を求めることで可能である．

#include <bits/stdc++.h>  
using namespace std;  
using ll = long long;  
using PII = pair<ll, ll>;  
#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)  
#define REP(i, n) FOR(i, 0, n)  
#define ALL(x) x.begin(), x.end()  
template<typename T> void chmin(T &a, const T &b) { a = min(a, b); }  
template<typename T> void chmax(T &a, const T &b) { a = max(a, b); }  
struct FastIO {FastIO() { cin.tie(0); ios::sync_with_stdio(0); }}fastiofastio;  
#ifdef DEBUG_   
#include "../program_contest_library/memo/dump.hpp"  
#else  
#define dump(...)  
#endif  
const ll INF = 1LL<<60;  
  
struct UnionFind {  
    vector<int> par, s;  
    UnionFind(int n=2e5) { init(n); }  
    void init(int n) {   
        s.assign(n, 1); par.resize(n);   
        iota(par.begin(), par.end(), 0);  
    }  
    int find(int x) {  
        if(par[x] == x) return x;  
        return par[x] = find(par[x]);  
    }  
    void unite(int x, int y) {  
        x = find(x);  
        y = find(y);  
        if(x == y) return;  
        if(s[x] < s[y]) par[x] = y, s[y] = s[x] + s[y];  
        else par[y] = x, s[x] = s[x] + s[y];  
    }  
    bool same(int x, int y) { return find(x) == find(y); }  
    int size(int x) { return s[find(x)]; }  
};  
  
struct LCA {  
    const int n = 0;  
    const int log2_n = 0;  
    vector<vector<int>> g;  
    vector<vector<int>> par;  // par[2^i上][頂点v]  
    vector<int> dep;  
  
    void dfs(int v, int p, int d) {  
        par[0][v] = p;  
        dep[v] = d;  
        for(auto to: g[v]) {  
            if(to == p) continue;  
            dfs(to, v, d+1);  
        }  
    }  
  
    LCA() {}  
    LCA(int n) : n(n), log2_n(log2(n)+1), g(n),  
        par(log2_n, vector<int>(n)), dep(n) {}  
  
    void add_edge(int u, int v) {  
        g[u].push_back(v);  
        g[v].push_back(u);  
    }  
  
    void build(ll root=0) {  
        dfs(root, -1, 0);  
        for(int i=0; i+1 < log2_n; ++i) {  
            for(int j = 0; j < n; ++j) {  
                if(par[i][j] < 0) par[i+1][j] = -1;  
                else par[i+1][j] = par[i][par[i][j]];  
            }  
        }  
    }  
  
    int get(int u, int v) {  
        if(dep[u] > dep[v]) swap(u, v);  
        REP(i, log2_n) {  
            if((dep[v] - dep[u]) >> i & 1) {  
                v = par[i][v];  
            }  
        }  
        if(u == v) return u;  
        for(int i=log2_n-1; i>=0; --i) {  
            if(par[i][u] != par[i][v]) {  
                u = par[i][u];  
                v = par[i][v];  
            }  
        }  
        return par[0][u];  
    }  
};  
  
int main(void) {  
    ll n, m, q;  
    cin >> n >> m >> q;  
    vector<vector<ll>> g(n);  
    UnionFind uf(n);  
    LCA graph(n+1);  
    REP(i, m) {  
        ll u, v;  
        cin >> u >> v;  
        u--, v--;  
        g[u].push_back(v);  
        g[v].push_back(u);  
        uf.unite(u, v);  
        graph.add_edge(u, v);  
    }  
    vector<ll> a(q), b(q), sz(n);  
    REP(i, q) {  
        cin >> a[i] >> b[i];  
        a[i]--, b[i]--;  
        // a[i] と b[i] が別の木に属するなら重心を求める対象に加える  
        if(!uf.same(a[i], b[i])) sz[a[i]]++, sz[b[i]]++;  
    }  
  
    // 各木に対して重心を求める  
    map<ll, ll> cs;  
    {   
        ll sum = 0;  
        function<void(ll,ll)> pre = [&](ll v, ll p) {  
            sum += sz[v];  
            for(auto to: g[v]) if(to != p) pre(to, v);  
        };  
        vector<bool> used(n);  
        vector<ll> centroid;  
        function<void(ll,ll)> dfs = [&](ll v, ll p) {  
            used[v] = true;  
            bool is_centroid = true;  
            for(auto to : g[v]) if (to != p) {  
                dfs(to, v);  
                if(sz[to] > sum/2) is_centroid = false;  
                sz[v] += sz[to];  
            }  
            if(sum-sz[v] > sum/2) is_centroid = false;  
            if(is_centroid) centroid.push_back(v);  
        };  
        REP(i, n) {  
            if(used[i]) continue;  
            sum = 0; centroid.clear();  
            pre(i, -1);  
            dfs(i, -1);  
            graph.add_edge(n, i); // LCAが木にしか対応してなかったので1頂点追加してそれを根にする  
            if(sum > 0) cs[uf.find(centroid[0])] = centroid[0];  
        }  
    }  
  
    graph.build(n);  
    ll ret = 0;  
    REP(i, q) {  
        // a[i] と b[i] が同じ木に属するなら2点の距離  
        if(uf.same(a[i], b[i])) {  
            ll lca = graph.get(a[i], b[i]);  
            ret += graph.dep[a[i]] + graph.dep[b[i]] - 2 * graph.dep[lca];  
        }  
        // 違う木に属するなら重心までの距離  
        else {  
            ll u = a[i], v = cs[uf.find(a[i])], lca = graph.get(u, v);  
            ret += graph.dep[u] + graph.dep[v] - 2 * graph.dep[lca];  
            u = b[i], v = cs[uf.find(b[i])], lca = graph.get(u, v);   
            ret += graph.dep[u] + graph.dep[v] - 2 * graph.dep[lca];  
        }  
    }  
    cout << ret << endl;  
  
    return 0;  
}