けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 235 E - MST + 1 (水色, 500 点)

MST の理解が問われる面白い教育的問題!

問題概要

頂点数  N、辺数  M の連結な重み付き単純無向グラフ  G が与えられる。なお、各辺の重みはすべて互いに異なることが保証されている。

 G の最小全域木を  T とする (一意に定まることが示せる)。

このグラフに  Q 個のクエリが与えられる。各クエリでは、2 頂点  u, v と重み  w が与えられる。この重み  w は、元のグラフ  G のどの辺の重みとも異なることが保証される。

このグラフ  G に頂点  u, v を結ぶ重み  w の辺  e を追加してできるグラフ  G' の最小全域木を  T' とする (一意に定まることが示せる)。

 T' が辺  e を含むかどうかを判定せよ。

制約

  •  2 \le N \le 2 \times 10^{5}
  •  N - 1 \le M \le 2 \times 10^{5}
  •  1 \le Q \le 2 \times 10^{5}

考えたこと

こういう問題大好き!

クエリ処理問題では、まず 1 つのクエリに対する問題が解けるかどうかを考えると良いことが多いと思う。

つまり、グラフ  G に対して 1 本の  e を追加することによって、最小全域木が変わるのかどうかを考えるのだ。

Kruskal 法を思い出そう。そうすると、「グラフ  G M 本の辺と  e を混ぜた  M+1 本の辺集合」に対して kruskal 法を適用したときに、辺  e を追加しようとする段階において


  •  e の両端点がすでに同じグループにあるとき: e はグラフ  G' の最小全域木にも採用されない (答えは "No")
  •  e の両端点が異なるグループにあるとき: e はグラフ  G' の最小全域木には採用される (答えは "Yes")

ということがわかる。

クエリが  Q 個あっても、結局すべてのクエリを先読みして、 M 本の辺に  Q 本の辺を混ぜた  M+Q 本の辺に対して Kruskal 法を適用することで解ける。

計算量は  O(N + (M + Q) \log (M + Q)) となる。

コード

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

// Union-Find
struct UnionFind {
    // core member
    vector<int> par;

    // constructor
    UnionFind() { }
    UnionFind(int n) : par(n, -1) { }
    void init(int n) { par.assign(n, -1); }
    
    // core methods
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool same(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 (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
    
    // debug
    friend ostream& operator << (ostream &s, UnionFind uf) {
        map<int, vector<int>> groups;
        for (int i = 0; i < uf.par.size(); ++i) {
            int r = uf.root(i);
            groups[r].push_back(i);
        }
        for (const auto &it : groups) {
            s << "group: ";
            for (auto v : it.second) s << v << " ";
            s << endl;
        }
        return s;
    }
};

struct Edge {
    int x, y, w, id;
    Edge() {}
    Edge(int x, int y, int w, int id) : x(x), y(y), w(w), id(id) {}
    friend bool operator < (const Edge &a, const Edge &b) { return a.w < b.w; }
};

int main() {
    int N, M, Q;
    cin >> N >> M >> Q;
    
    vector<Edge> edges;
    for (int i = 0; i < M + Q; ++i) {
        int a, b, c;
        cin >> a >> b >> c;
        --a, --b;
        edges.push_back(Edge(a, b, c, i));
    }
    sort(edges.begin(), edges.end());
    
    UnionFind uf(N);
    vector<bool> res(Q);
    for (const auto &e : edges) {
        if (e.id >= M) res[e.id - M] = !uf.same(e.x, e.y);
        else uf.merge(e.x, e.y);
    }
    
    for (int i = 0; i < Q; ++i) cout << (res[i] ? "Yes" : "No") << endl;
}