CS Academy 082 DIV1 E - Water Supply - けんちょんの競プロ精進記録

ラグランジュ緩和か、、、しかしこれ、ある特定の 1 頂点について次数が $k$ 以下であるような最小全域木を求める問題とみなせるわけだけど、こんなんが解けるのは面白い。

全域最小木スライドにある通り、全頂点について次数が $k$ 以下となるような最小全域木を求める問題は NP-hard。

問題へのリンク

問題概要 (意訳)

$N+1$ 頂点からなる重みつき単純無向密グラフが与えられる。 $N+1$ 個の頂点のうちの 1 点 $s$ が特別に指定されている。

$s$ における次数が $K$ 以下となるような全域木のうち、重みの合計の最小値を求めよ。

制約

$1 \le (テストケース数 T) \le 2000$
$1 \le K \le N \le 2000$

考えたこと

satanic にゃんの解説にすべてが書いてある。

アイディアとして

普通に $s$ での次数を制限せずに MST を求めると、 $s$ での次数が $K$ を超えてしまう可能性がある (超えなければそれがすでに答え)
$s$ から出ている辺の重みを一律に v だけ増やすと、 $s$ の次数が減少した解が得られる
$v$ を二分探索しながら調節する

正当性などは、satanic さんの解説を参照。で、注意点として

毎回二分探索していては $O(TN^{2}\log{N}\log{C})$ かかってしまって間に合わない ( $C$ は重み最大値)

そこで、最初に $s$ を除いて最小全域木を求めてみる。このときに使われない辺は、どのみち使うことは絶対にない。よって、毎回の二分探索で使用する辺を $O(N)$ 個に絞ることができる (高難易度でよく見る考え方なイメージ)！！！

こうすると、計算時間は $O(TN\log{N}\log{C})$ に減って間に合う。

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

struct UnionFind {
    vector<int> par;
    vector<int> rank;
    
    UnionFind(int n) { init(n); }
    void init(int n) {
      par.resize(n); rank.resize(n);
      for (int i = 0; i < n; ++i) par[i] = i, rank[i] = 0;
    }
    int root(int x) {
        if (par[x] == x) return x;
        else return par[x] = root(par[x]);
    }
    bool issame(int x, int y) return root(x) == root(y);
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (rank[x] < rank[y]) swap(x, y);
        if (rank[x] == rank[y]) ++rank[x];
        par[y] = x;
        return true;
    }
};

int TT;
int N, K;
vector<long long> F;
vector<vector<long long> > G;

typedef pair<int,int> pint;
typedef pair<long long, pint> Edge;

bool cmp(const Edge &a, const Edge &b) {
  return (a.first < b.first) || (a.first == b.first && a.second > b.second);
}

int main() {
    cin >> TT;
    for (int CASE = 0; CASE < TT; ++CASE) {
        cin >> N >> K;
        F.resize(N); for (int i = 0; i < N; ++i) cin >> F[i];
        vector<Edge> alledges;
        for (int i = 0; i < N; ++i) for (int j = 0; j < N; ++j) {
            long long cost; cin >> cost;
            if (i < j) alledges.push_back(Edge(cost, pint(i, j)));
          }
        sort(alledges.begin(), alledges.end());

        // 考慮すべき枝を絞る: fortune なしのクラスカル法で選ばれる枝以外は考えなくて良い
        vector<Edge> edges;
        UnionFind uf(N);
        for (auto e : alledges) {
          int u = e.second.first, v = e.second.second;
          if (uf.issame(u, v)) continue;
          uf.merge(u, v);
          edges.push_back(e);
        }

        // ラグランジュ緩和しつつ、各ステップは絞られた枝 + fortune で Kruskal
        long long res = 1LL<<60;
        long long low = -1, high = 1LL<<40;
        while (high - low > 1) {
            long long mid = (high + low) / 2;
            vector<Edge> lagedges = edges;
            for (int i = 0; i < N; ++i) lagedges.push_back(Edge(F[i] + mid, pint(N, i)));
            sort(lagedges.begin(), lagedges.end(), cmp);

            long long allcost = 0;
            int num_connected = 0;
            int num_fortune = 0;
            UnionFind uf(N+1);
            bool over = false; // 
            for (auto e : lagedges) {
              int u = e.second.first, v = e.second.second;
              if (uf.issame(u, v)) continue;
              if (u == N) {
                if (num_fortune >= K) {
                  over = true;
                  continue;
                }
                ++num_fortune;
              }
              uf.merge(u, v);
              allcost += e.first;
              ++num_connected;
            }

            if (num_connected == N) res = min(res, allcost - mid * num_fortune);
            if (!over) high = mid;
            else low = mid;
        }
        cout << res << endl;
    }
}