AOJ 2559 Minimum Spanning Tree

AOJ 競プロ

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2559

解法

マージテクだけで解ける．
最初に MST を構成しておく．MST に含まれない辺は自明．
以下MSTに含まれる辺についてコストを求める．
MST 上を DFS で探索し，各部分木から外に生えていく辺を全部求めていく．
外に出ていく辺としては以下の2通りある．

  • 行き先は別の部分木の頂点（先祖ー子孫の関係にない）
  • 行き先が祖先

どちらの場合も同じやり方で処理できる．

DFS の戻り値として，その部分木から外に出ていくすべての辺を返すことにする．
DFS の過程で今頂点 v にいるとすると，

  1. 辺集合の格納先を res とする．
  2. 各部分木に対して辺集合を求め，マージテクで res と併合する．このとき，intersect する場合は削除するようにし，そうでない場合は挿入するようにする．
  3. v から出ている（MST上にない）すべての辺 e に対して，e が res にあれば削除，なければ挿入する．
  4. res を返す

ステップ2が「行き先が別の部分木の頂点の辺」を打ち消す操作，ステップ3が「行き先が祖先の辺」を打ち消す操作に対応する．
res を set で管理すれば，マージテクしながら削除にも対応できる．このとき，辺の順序は (cost, idx) で持っておけば良い．

計算量は O(n + m(logm)^2) である．

ソースコード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

class union_find {
public:
    union_find(int n) : par(n, -1) {}

    int root(int x) {
        return par[x] < 0 ? x : par[x] = root(par[x]);
    }

    bool unite(int x, int y) {
        x = root(x), y = root(y);
        if(x == y) return false;
        if(par[x] > par[y]) swap(x, y);
        par[x] += par[y];
        par[y] = x;
        return true;
    }

private:
    vector<int> par;
};

struct edge {
    int from, to, cost, idx;
    bool operator<(const edge& e) const {
        return make_pair(cost, idx) < make_pair(e.cost, e.idx);
    }
};

int main() {
    ios_base::sync_with_stdio(false); cin.tie(0);

    int n, m; cin >> n >> m;
    vector<vector<edge>> g(n);
    vector<edge> es;
    for(int i = 0; i < m; ++i) {
        int a, b, w; cin >> a >> b >> w;
        a--, b--;
        g[a].push_back(edge{a, b, w, i});
        g[b].push_back(edge{b, a, w, i});
        es.push_back(edge{a, b, w, i});
    }
    sort(es.begin(), es.end());

    vector<vector<edge>> mst(n);
    vector<bool> in_mst(m);
    vector<ll> ans(m, -1);
    ll mst_cost = 0;
    {
        union_find uf(n);
        for(auto const& e : es) {
            if(uf.unite(e.from, e.to)) {
                in_mst[e.idx] = true;
                mst_cost += e.cost;
                mst[e.from].push_back(edge{e.from, e.to, e.cost, e.idx});
                mst[e.to].push_back(edge{e.to, e.from, e.cost, e.idx});
            }
        }
        for(int i = 0; i < m; ++i) {
            if(!in_mst[i]) {
                ans[i] = mst_cost;
            }
        }
    }

    function<set<edge>(int, int)> solve = [&] (int v, int p) {
        set<edge> res;
        vector<set<edge>> buf_es;
        for(auto const& e : mst[v]) {
            if(e.to == p) continue;
            auto tmp_es = solve(e.to, v);
            tmp_es.erase(e);
            if(!tmp_es.empty()) {
                ans[e.idx] = mst_cost - e.cost + tmp_es.begin()->cost;
            }
            if(res.size() < tmp_es.size()) swap(res, tmp_es);
            for(auto const& e : tmp_es) {
                if(res.count(e)) res.erase(e);
                else             res.insert(e);
            }
        }
        for(auto const& e : g[v]) {
            if(in_mst[e.idx]) continue;
            if(res.count(e) == 1) res.erase(e);
            else                  res.insert(e);
        }
        return res;
    };
    solve(0, -1);

    for(auto const x : ans) {
        cout << x << "\n";
    }
}

感想

重心分解の例題で見かけたので復習がてら解いたのですが，重心分解いらないじゃんとなってしまい悲しい．