ゆらのふなびと

競プロ, Python, C++

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

問題

N個の都市がM本の道路でつながれている。i本目の道路は都市a_iとb_iを結び、y_i年に造られたものである。

Q人の国民がいる。国民jは都市v_jに住んでおり、作られた都市がw_jより大きな道路しか使わない。

各国民について、自身が住んでいる都市から条件を満たす道路のみを用いて到達可能な都市の個数を求めよ。

制約

  • 1≦N≦100,000
  • 0≦M≦200,000
  • 1≦ai,bi≦N
  • ai≠bi
  • 1≦yi≦200,000
  • 1≦Q≦100,000
  • 1≦vj≦N
  • 0≦wj≦200,000

解法

問題を言い換えると、国民jについての答えは「w_jより大きな重みの辺のみを使ったとき、v_iが属する連結成分には何個の頂点が含まれるか?」である。連結成分の管理はUnionFindで簡単にできるが、国民ごとに毎回w_jより大きな重みの辺を考えて再構築していたのではO(MQ)*1かかり間に合わない。

そこで、辺とクエリを年の降順にソートする。ソートされた各クエリごとに、そのクエリの年より大きな年の辺だけunion→個数を記録 を繰り返し、最後に並べ直して出力する。これでO(max(MlogM, QlogQ))で解けた。

テクニックとしてはこれ↓に似ている。

arc042.contest.atcoder.jp

#include <bits/stdc++.h>
using namespace std;
#define int long long   // <-----!!!!!!!!!!!!!!!!!!!

#define rep(i,n) for (int i=0;i<(n);i++)
#define rep2(i,a,b) for (int i=(a);i<(b);i++)
#define rrep(i,n) for (int i=(n)-1;i>=0;i--)
#define rrep2(i,a,b) for (int i=(b)-1;i>=(a);i--)
#define all(a) (a).begin(),(a).end()

typedef long long ll;
typedef pair<int, int> P;
typedef tuple<int, int, int> TUPLE;
typedef vector<int> V;
typedef vector<V> VV;
typedef vector<VV> VVV;

class UnionFind {
private:
    int n;
    vector<int> uni;
public:
    UnionFind(int _n) {
        n = _n;
        uni.clear();
        uni.resize(n, -1);
    }
    int root(int x) {
        if (uni[x] < 0) return x;
        return uni[x] = root(uni[x]);
    }
    bool same(int x, int y) {
        return root(x) == root(y);
    }
    void unite(int x, int y) {
        x = root(x);
        y = root(y);
        if (x == y) return;
        if (uni[x] > uni[y]) swap(x, y);
        uni[x] += uni[y];
        uni[y] = x;

    }
    int getSize(int x) {
        return -uni[root(x)];
    }
    void print() {
        for (auto x : uni) cout << x << " ";
        cout << endl;
    }
};

signed main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(0);

    int N, M;
    cin >> N >> M;
    VV edges;
    rep(i, M) {
        int y, a, b;
        cin >> a >> b >> y;
        a--, b--;
        edges.emplace_back(V{y, a, b});
    }
    sort(all(edges));
    reverse(all(edges));

    int Q;
    cin >> Q;
    VV queries;
    rep(i, Q) {
        int a, y;
        cin >> a >> y;
        a--;
        queries.emplace_back(V{y, a, i});
    }
    sort(all(queries));
    reverse(all(queries));

    UnionFind uf(N);
    map<int, int> ans;
    int i = 0;
    for (auto q : queries) {
        while (i < M && q[0] < edges[i][0]) {
            uf.unite(edges[i][1], edges[i][2]);
            i++;
        }
        ans[q[2]] = uf.getSize(q[1]);
    }

    for (auto &p : ans) {
        cout << p.second << endl;
    }
}

UnionFindのライブラリにgetSizeが仲間入りした。

*1:アッカーマンのアレは省略