CODE FESTIVAL 2017 Elimination Tournament Round 2 A - Colorful MST

あさぷろで解いたときは解けなかった。面白い問題。
A - Colorful MST

問題

N頂点M辺のグラフがある。各頂点には色1から色Kまでのどれかの色が塗られているか、あるいはまだ色が塗られていない。
以下の操作を行ったときにできるグラフの最小全域木の重みの最小値を求めよ。どのようにしても最小全域木を作れない場合は-1を出力せよ。

  1. まだ色が塗られていない頂点に色1から色Kまでのどれかの色を塗る。別の頂点には別の色を塗ってよい。
  2. 同じ色の頂点を1点に潰して、K頂点のグラフを作る。このとき、対応する色の頂点が存在しないような色については、その頂点は0本の辺が接続することになる。

制約

  • 1 <= K <= N <= 10^5
  • 1 <= M <= 10^5
  • 1 <= (辺の重み) <= 10^9
  • 与えられるグラフは、単純とも連結とも限らない

解法

適当に(K - 1)辺(実行可能解)を取った時、その辺集合が全域木となるような色の塗り方は存在するだろうか? 存在するためには以下の条件が必要十分:

  • すでに塗られている頂点を色ごとに1点に潰した時、潰した後の点集合が森になっている。

(必要性): どのような塗り方をしても、塗られていない状態よりは潰れるので、森になりにくい。よって自明
(十分性): 塗り方を作ればよい。まだ塗られていない頂点に接続する辺1つにつき、まだ使われていない色を1つずつ使っていけばよい。
以上より、すでに塗られている頂点を同一視しながらKruskal法を実行するだけで解ける。

登場する典型

  • 実行可能解の条件を考える

実装上の注意点

  • オーバーフローの危険があるのでlong longを使う
  • ちょうど(K - 1)個の辺を選ぶ
    • 1敗 (は?)

提出: Submission #4343453 - CODE FESTIVAL 2017 Elimination Tournament Round 2 (Parallel) (C++)

const int N = 112222;
int a[N], b[N], w[N];

typedef pair<ll, PI> PLPI;

int main(void) {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, m, k;
  cin >> n >> m >> k;
  VI c(n);
  REP(i, 0, n) {
    cin >> c[i];
  }
  vector<PLPI> pool;
  REP(i, 0, m) {
    cin >> a[i] >> b[i] >> w[i];
    a[i]--, b[i]--;
    pool.push_back(PLPI(w[i], PI(a[i], b[i])));
  }
  sort(pool.begin(), pool.end());
  UnionFind uf(k + n);
  int conn = k;
  ll tot = 0;
  REP(i, 0, pool.size()) {
    if (conn <= 1) break;
    ll ww = pool[i].first;
    int x = pool[i].second.first;
    int y = pool[i].second.second;
    x = c[x] == 0 ? k + x : c[x] - 1;
    y = c[y] == 0 ? k + y : c[y] - 1;
    assert (x >= 0);
    assert (y >= 0);
    if (not uf.is_same_set(x, y)) {
      tot += ww;
      uf.unite(x, y);
      conn--;
    }
  }
  cout << (conn > 1 ? -1 : tot) << "\n";
}

まとめ

実行可能解の条件を考えるのも割とよく見る典型という気がする。