(拡張UnionFind・ラムダ式)ABC-040-D-道路の老朽化対策について

問題URL→http://abc040.contest.atcoder.jp/tasks/abc040_d

クエリを年代順にソートしてUnionFindする問題。
ただし普通のUnionFindではだめで、各集合のサイズも取得できるようにしないといけない。
今回はそれを実装した。
またUnionFindを汎用的に使えるようにclass化もしたので良いライブラリが出来たと思う。

ラムダ式便利。

#include <iostream>
#include <vector>
#include <string.h>
#include <climits>
#include <algorithm>
#define FOR(i,bg,ed) for(ll i=(bg);i<(ed);i++)
#define REP(i,n) FOR(i,0,n)
#define MOD 1000000007
#define int long long
#ifdef int
const long long INF = LLONG_MAX / 10;
#else
const int INF = 1010101010;
#endif
using namespace std;
typedef long long ll;
typedef vector<vector<ll>> mat;
typedef pair<int, int> P;
//typedef pair<double, double> P;

//UnionFind!!!
class UnionFind
{
private:
    vector<int> par;
    vector<int> myRank;
    vector<int> cnt;

public:
    //n要素で初期化
    UnionFind(int n)
    {
        par.resize(n);
        myRank.resize(n);
        cnt.resize(n);
        REP(i,n) {
            par[i] = i;
            myRank[i] = 0;
            cnt[i] = 1;
        }
    }

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

    //xとyの属する集合を併合
    void unite(int x, int y)
    {
        x = find(x);
        y = find(y);
        if (x == y) return;

        if (myRank[x] < myRank[y]) {
            par[x] = y;
            cnt[y] += cnt[x];
        } else {
            par[y] = x;
            cnt[x] += cnt[y];
            if (myRank[x] == myRank[y]) myRank[x]++;
        }
    }

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

    //xが含まれる集合の要素数
    int count(int x)
    {
        return cnt[find(x)];
    }
};

int N, M;
int a[202020], b[202020], y[202020];
int Q;
int v[101010], w[101010];

struct Query
{
    int year;
    int qType;
    int idx;
    int a, b;
    int v;
};

Query qs[303030];

signed main()
{
    cin >> N >> M;
    REP(i,M) {
        cin >> qs[i].a >> qs[i].b >> qs[i].year;
        qs[i].a--; qs[i].b--;
        qs[i].qType = 1;
    }
    cin >> Q;
    REP(i,Q) {
        cin >> qs[M+i].v >> qs[M+i].year;
        qs[M+i].v--;
        qs[M+i].qType = 2;
        qs[M+i].idx = i;
    }

    sort(qs, qs + M + Q, [](Query a, Query b){
        if (a.year == b.year) {
            return a.qType > b.qType;
        } else {
            return a.year > b.year;
        }
    });

    /*
    REP(i,M+Q) {
        cout << qs[i].year << " : " << qs[i].qType << endl;
    }
    */

    UnionFind uf(101010);

    vector<int> ans(Q);

    REP(i,M+Q) {
        if (qs[i].qType == 1) {
            uf.unite(qs[i].a, qs[i].b);
        } else {
            ans[qs[i].idx] = uf.count(qs[i].v);
        }
    }

    REP(i,Q) {
        cout << ans[i] << endl;
    }
}