第二回アルゴリズム実技検定 (PAST) O - 可変全域木

問題: O - Variable Spanning Trees

問題概要

$N ( \leq 10^5)$ 頂点 $M (\leq 10^5)$ 辺の無向グラフが与えられる。$i$ 番目の辺は頂点 $a_i, b_i$ を結び，$c_i$ の重みを持つ。

各辺について，その辺を含む最小の重みの全域木を求め，その重みを出力せよ。

解法概要

「各辺について，その辺を含む最小の重みの全域木」の太字の部分さえなければ，ただの最小全域木問題となる。その場合，UnionFind を使ったクラスカル法などをただ適用するだけで良い。今回の問題でも，注目する辺が最小全域木に初めから含まれる場合はクラスカル法を適用するだけだ。

考えなければならないのは，注目する辺が最小全域木に含まれない場合である。注目する辺を最初に使うようにクラスカル法をアレンジしようとしても，時間計算量 $\mathcal{O}(M^2 \log M)$ (クラスカル法は $\mathcal{O}(M \log M)$ なので) となり TLE してしまう。なんとかして1辺あたり対数時間や定数時間で求めたい。

そこで，最小全域木と「注目する辺を含む最小の重みの全域木」の差分を考えてみる。以下のような全域木があった場合に，辺 (2, 4) を追加で使うことを考えるとしよう。

(画像: GRAPH × GRAPH：競技プログラミングにおけるグラフ可視化サイト のスクショ)

もともと木なので，当然辺 (2, 4) を追加したら閉路ができる。辺 (2, 4) は必ず使うことにすると，他を1辺削除しないといけない。重みを最小にしたいので， (2, 4) を含む閉路 2-1-3-4 の中で，辺 (4, 5) 以外でもっとも重みの大きい辺を取り除くことになる。同様に，辺 (4, 5) を追加することを考える場合には，閉路 4-3-5 の中で (4, 5) 以外でもっとも重みの大きい辺を取り除くことになる。

閉路についてもう少し考えてみると，辺 (a, b) を加えてできる閉路は a - ... - (根からもっとも遠い，a と b の共通祖先 (LCA)) - ... - b のようになっていることがわかる。したがって，この閉路に含まれる辺 (a, b) 以外の辺でもっとも重みの大きい辺の重みは，a と b の LCA を l とすると「a から l までのパスに含まれる辺の最大の重み」と「b から l までのパスに含まれる辺の最大の重み」のうちより大きいほうである。

これは LCA をダブリングで求めるのと同じ要領で求めることができる。つまり，ダブリングによって LCA を求めるときに各頂点の $2^j$ 個上の頂点を持っておくように，頂点 $i$ から頂点 $i$ の $2^j$ 個上の頂点に移動するまでに通る辺の中での最大の重みを持っておけば良い。

ダブリングによって LCA と LCA からその頂点までのパス中の最大重みを求めるのは，前処理 $\mathcal{O}(N \log N)$，LCA・最大重み計算 $\mathcal{O}(\log N)$ でできる。

よって全体での時間計算量が $\mathcal{O}(M \log M + M \log N)$ となり，間に合う。

ソースコード

#include <iostream>
#include <vector>
#include <algorithm>
#include <tuple>

using namespace std;

using ll =  long long;
using Pll = pair<ll, ll>;
using Pii = pair<int, int>;

constexpr int N_MAX = 100000;

struct Edge {
    int i, from, to;
    ll cost;
    Edge(int i, int from, int to, ll cost): i(i), from(from), to(to), cost(cost) {};
    bool operator < (const Edge& e) const {
        return cost < e.cost;
    }
};

vector<tuple<int, ll>> graph[N_MAX];
vector<Edge> edges;

class UnionFind {
    public:
    int n;
    vector<int> par, rank;

    UnionFind(int n) {
        for(int i=0;i<n;++i) {
            par.push_back(i);
            rank.push_back(0);
        }
    }

    int find(int x) {
        return ((par[x] == x) ? x : par[x] = find(par[x]));
    }

    void unite(int x, int y) {
        x = find(x);
        y = find(y);

        if(x == y) return;

        if(rank[x] < rank[y]){
            par[x] = y;
        }else{
            par[y] = x;
            if(rank[x] == rank[y]) rank[x]++;
        }
        return;
    }

    bool is_same(int x, int y) {
        return find(x) == find(y);
    }
};

class LowestCommonAncestor {
    public:
    int root, n, LOG_V_MAX = 30;
    vector<int> depth;
    vector<vector<int>> parent;
    vector<vector<ll>> max_cost;    // 頂点 i から 2^j 個上のノードまでの辺のうちコストが最大のもののコスト

    LowestCommonAncestor(int root=0, int n=N_MAX): root(root), n(n) {
        depth.assign(n, -1);
        parent.assign(n, vector<int>(LOG_V_MAX, -1));
        max_cost.assign(n, vector<ll>(LOG_V_MAX, 0));
        dfs(root, -1, 0);
        for(int j=1;j<LOG_V_MAX;++j) {
            for(int i=0;i<n;++i) {
                if(parent[i][j-1] == -1) continue;
                parent[i][j] = parent[parent[i][j-1]][j-1];
                max_cost[i][j] = max(max_cost[i][j-1], max_cost[parent[i][j-1]][j-1]);
            }
        }
    }

    void dfs(int node, int par, int d) {
        depth[node] = d;
        parent[node][0] = par;
        for(auto child: graph[node]) {
            int c = get<0>(child);
            if(c == par) continue;
            max_cost[c][0] = get<1>(child);
            cerr << max_cost[c][0] << endl;
            dfs(c, node, d+1);
        }
    }

    ll get_max_cost(int u, int v) {
        if(depth[u] > depth[v]) swap(u, v);
        ll ret = 0LL;
        int depth_diff = depth[v] - depth[u];
        for(int j=0;j<LOG_V_MAX;++j) {
            if(depth_diff & (1 << j)) {
                ret = max(ret, max_cost[v][j]);
                v = parent[v][j];
            }
        }
        if(u == v) return ret;
        for(int j=LOG_V_MAX-1;j>=0;--j) {
            if(parent[u][j] != parent[v][j]) {
                ret = max({ret, max_cost[u][j], max_cost[v][j]});
                u = parent[u][j];
                v = parent[v][j];
            }
        }
        ret = max({ret, max_cost[u][0], max_cost[v][0]});
        return ret;
    }
    
};


int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr);
    int n, m, a, b; ll c;
    cin >> n >> m;
    for(int i=0;i<m;++i) {
        cin >> a >> b >> c;
        --a; --b;
        edges.emplace_back(i, a, b, c);
    }
    sort(edges.begin(), edges.end());

    UnionFind uf = UnionFind(n);
    ll mst_cost = 0LL;
    vector<bool> used(m, false);
    for(int i=0;i<m;++i) {
        if(uf.is_same(edges[i].from, edges[i].to)) continue;
        uf.unite(edges[i].from, edges[i].to);
        mst_cost += edges[i].cost;
        graph[edges[i].from].emplace_back(edges[i].to, edges[i].cost);
        graph[edges[i].to].emplace_back(edges[i].from, edges[i].cost);
        used[i] = true;
    }

    vector<ll> ans(m);
    LowestCommonAncestor lca = LowestCommonAncestor(0, n);
    for(int i=0;i<m;++i) {
        if(used[i]) {
            ans[edges[i].i] = mst_cost;
        } else {
            ans[edges[i].i] = mst_cost + edges[i].cost - lca.get_max_cost(edges[i].from, edges[i].to);
        }
    }

    for(int i=0;i<m;++i) cout << ans[i] << endl;
}

所感

えびちゃんの日記 を見る限り，こどふぉで既出らしい。それと，HL 分解を持っていると貼るだけになるらしい。HL 分解は HL 分解でないと解けない問題になかなか出くわさないため勉強モチベが上がらないなぁ……。

公式解説 曰く，第一の想定解はダブリングでもなく HL 分解でもなくオイラーツアー でもなく，並列二分探索のようだけど。せめて並列二分探索解は理解するか。