ひらめの日常

プログラミングと読書と

ABC040-D 道路の老朽化対策について (500)

道路の老朽化対策について

タイトルだけ見ると政策の説明みたいだけど、競技プログラミングの問題。

問題はこちら。 atcoder.jp

都市をつなぐ道路が与えられるので、都市を頂点、道路を辺としてみたグラフを考えることにする。

Unionfindを使ってクエリごとに構築して、大きさを調べようと思ったが、これだとTLEするので工夫が求められる。

考え方

1クエリごとに木を再構築していると、 O(QM) かかるので間に合わない。そこで、クエリを先に読んでおいて、 wが大きい順にクエリを処理する。こうすることで、既に存在する木に新たに条件を満たす辺を加えていくだけでよくなる。

計算量は、各辺を見る回数が高々1回なので、操作回数が O(Q + M)になる。

解答

答えを出力するときには元の順番で出力しなければいけないので、クエリのindexを保持しておく必要がある。

また、sortがしやすくなるようにpairの順番を指定している。

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
#define rep(i, n) for(ll i = 0; i < (ll)(n); i++)

/* ------------- ANSWER ------------- */
/* ---------------------------------- */
class UnionFind {
private:
    vector<ll> size; // グループに属する物の数.
public:
    vector<ll> par; // 親
    vector<ll> rank; // 木の深さ

    explicit UnionFind(unsigned int n) {
        par.resize(n);
        rank.resize(n);
        size.resize(n);
        rep(i, n) {
            par[i] = i;
            rank[i] = 0;
            size[i] = 1;
        }
    }

    // 木の根を求める
    ll find(ll x) {
        if (par[x] == x) {
            return x;
        } else {
            return par[x] = find(par[x]);
        }
    }

    // グループのサイズを求める.
    ll calc_size(ll x) {
        return size[find(x)];
    }

    // xとyの属する集合を併合
    void unite(ll x, ll 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]++;
        }
        size[x] = size[y] = size[x] + size[y];
    }

    // xとyが同じ集合に属するか否か
    bool is_same(ll x, ll y) {
        return find(x) == find(y);
    }
};


int main() {
    ll n, m;
    cin >> n >> m;
    // yab[i].first: y
    // yab[i].second.first: a
    // yab[i].second.second: b
    priority_queue<pair<ll, pair<ll, ll>>> yab;
    rep(i, m) {
        ll a, b, y;
        cin >> a >> b >> y;
        a--, b--;
        yab.push({y, {a, b}});
    }

    ll q;
    cin >> q;
    // wvi[i].first.first: w
    // wvi[i].first.second: v
    // wvi[i].second: index
    vector<pair<pair<ll, ll>, ll>> wvi(q);
    rep(i, q) {
        ll v, w;
        cin >> v >> w;
        v--;
        wvi[i] = {{w, v}, i};
    }
    sort(wvi.begin(), wvi.end(), greater<>());

    UnionFind uf(n);

    vector<ll> ans(q);
    rep(i, q) {
        while (!yab.empty() && yab.top().first > wvi[i].first.first) {
            pair<ll, ll> ab = yab.top().second;
            uf.unite(ab.first, ab.second);
            yab.pop();
        }
        ans[wvi[i].second] = uf.calc_size(wvi[i].first.second);
    }
    rep(i, q) cout << ans[i] << endl;
    return 0;
}

Submission #6310485 - AtCoder Beginner Contest 040