AtCoder Regular Contest 029 C - 高橋君と国家

問題

　 $N$ 個の都市と $M$ 本の道を持つ国家がある。都市には交易所を置くという操作ができ、道には舗装するという操作ができ、それぞれ各都市、道ごとにコストが設定されている。全ての都市について、「その都市に交易所が設置されている」あるいは「舗装された道のみを用いて交易所が設置されている都市に移動できる」条件を満たすように操作を行うとき、最小のコストを求めよ。

解法

　 $0$ 番目の都市というものを考え、ここにはすでに交易所があり他の全ての都市へ道を持つものとする。「都市 $i$ に交易所を置く」という操作は「都市 $0$ と都市 $i$ の道を舗装する」という操作に変換して考えるものとする。このコストはもともと都市 $i$ に交易所を設置するコストとして与えられたものに等しい。

　このとき、条件を満たす状態は都市 $0$ から都市 $N$ が連結になっていることであり、これはプリム法やクラスカル法などから高速に求めることができる。

反省

　1時間11分(4WA)でAC。1時間ほどダイクストラ法 + Union Findみたいな方針でやっていたが、WAがしこたま出てどうしようもない感じだったので解説スライドを見た。全く想定していない解法だった。

　フローの問題とかである(らしい)感じの、仮想的なノードを作って問題を簡単にする手法にはまだ全然慣れていなくてどういうときに有効なのかよくわからない。しかしより一般化すると、「操作が複数あるときにそれを統一的な視点から眺め直す」ということなのだろうか。交易所の設置と道の舗装という操作を一つにまとめたいという気になって、こういうことが思いつくのかもしれない。交易所の設置コストは完全に初期値としか思っていなかった。

　実装が遅くていろいろ考察を試せなかったのも反省点。そもそもふわっとした考察のまま適当に実装し始めてしまったのも良くなかった。

コード

#include"bits/stdc++.h"
using namespace std;
using ll = int64_t;

struct Edge {
    ll from, to, cost;
};
vector<ll> parent;

ll root(ll x) {
    return (parent[x] == x ? x : parent[x] = root(parent[x]));
}

void unite(ll x, ll y) {
    x = root(x);
    y = root(y);
    if (x != y) {
        parent[y] = x;
    }
}

bool same(ll x, ll y) {
    return (root(x) == root(y));
}

int main() {
    ll N, M;
    cin >> N >> M;
    vector<ll> c(N);
    for (ll i = 0; i < N; i++) {
        cin >> c[i];
    }
    vector<Edge> edges(N);
    for (ll i = 0; i < M; i++) {
        ll a, b, r;
        cin >> a >> b >> r;
        a--; b--;
        edges.push_back({ a, b, r });
    }
    for (ll i = 0; i < N; i++) {
        edges.push_back({ i, N, c[i] });
    }

    parent.resize(N + 1);
    iota(parent.begin(), parent.end(), (ll)0);

    sort(edges.begin(), edges.end(), [&](const Edge& lhs, const Edge& rhs) {
        return lhs.cost < rhs.cost;
    });

    ll ans = 0;
    for (Edge e : edges) {
        if (!same(e.from, e.to)) {
            ans += e.cost;
            unite(e.from, e.to);
        }
    }

    cout << ans << endl;
}