UPC10525: Dove 打扑克

#include "bits/stdc++.h"

using namespace std;
typedef long long ll;
const int maxn = 3e5 + 10;
ll fa[maxn], cnt[maxn], val[maxn];
ll all;
set<ll> num;

ll findd(ll x) {
    return fa[x] == x ? x : fa[x] = findd(fa[x]);
}

void uni(ll a, ll b) {
    ll p1 = findd(a), p2 = findd(b);
    if (p1 != p2) {
        cnt[val[p1]]--;
        cnt[val[p2]]--;
        if (cnt[val[p1]] == 0) num.erase(val[p1]);
        if (cnt[val[p2]] == 0) num.erase(val[p2]);
        cnt[val[p1] + val[p2]]++;
        num.insert(val[p1] + val[p2]);
        fa[p2] = p1;
        val[p1] += val[p2];
        all--;
    }
}

void init(ll n) {
    cnt[1] = n;
    all = n;
    num.insert(1);
    for (ll i = 1; i <= n; i++) {
        fa[i] = i;
        val[i] = 1;
    }
}

int main() {
    freopen("input.txt", "r", stdin);
    ll n, m;
    scanf("%lld %lld", &n, &m);
    init(n);
    ll op;
    ll x, y;
    for (ll i = 0; i < m; i++) {
        scanf("%lld", &op);
        if (op == 1) {
            scanf("%lld %lld", &x, &y);
            uni(x, y);
        } else {
            scanf("%lld", &x);
            set<ll>::iterator it, j;
            it = num.begin();
            j = it;
            ll sum = 0;
            ll ans = 0;
            for (; it != num.end(); it++) {
                while (j != num.end() && (*j - *it < x || *j <= *it)) {
                    sum += cnt[*j];
                    j++;
                }
                sum -= cnt[*it];
                if (x > 0) ans += 1ll * cnt[*it] * (cnt[*it] - 1) / 2;
                ans = ans + cnt[*it] * sum;
            }
            printf("%lld\n", all * (all - 1) / 2 - ans);
        }
    }
    return 0;
}