HackerRank Journey to the Moon

気持ちが良い問題だった
問題リンク

問題概要

N人の宇宙飛行士がいます
P個のペアが渡されます
ペアは同じ国籍同士の宇宙飛行士の番号です
N人の中から，国籍が違う宇宙飛行士のペアはいくつあるか

解説

同じ国籍どうしの宇宙飛行士をDFSやUnionFindでグループを作ります
例えばAグループに3人，Bグループに4人いるとすると，Aグループ3人それぞれに対しBグループの4人がペアになりうるので生成できるペア数は12となります

愚直に二重for分でペア数を掛けていくと，時間がかかるのは予想できます

そこで，ひとまずN人から2人を選ぶ組み合わせの総数を出します
そこから同じグループから2人を選ぶ組み合わせの数をグループごとに見ていき引いていけば，答えが求まります

イメージは余事象的な感じですかね？

コード(C++14)

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

struct UnionFind{
    vector<long long> par;
    UnionFind(long long N) : par(N){
        for(long long i = 0; i < N; i++){
            par[i] = -1;
        }
    }
    long long root(long long x){
        if(par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    void unite(long long x, long long y){
        long long rx = root(x);
        long long ry = root(y);
        if(rx == ry) return;
        else{
            if(par[rx] > par[ry]) swap(rx, ry);
            par[rx] += par[ry];
            par[ry] = rx;
       }
    }
    bool same(long long x, long long y){
        return root(x) == root(y);
    }
    void check(){
        for(auto x : par){
            cout << x << " ";
        }
        cout << endl;
    }
};

int main(){
    cin.tie(nullptr);
    ios::sync_with_stdio(false);

    long long n, p;
    cin >> n >> p;
    UnionFind uf(n);
    while(p--){
        long long a, b;
        cin >> a >> b;
        uf.unite(a, b);
    }

    map<long long, long long> mp;
    for(long long i = 0; i < n; i++){
        long long r = uf.root(i);
        mp[r] = -uf.par[r];
    }

    vector<long long> B;
    for(auto x : mp){
        B.push_back(x.second);
    }

    long long bs = B.size();
    long long ans = n * (n - 1) / 2;
    for(long long i = 0; i < bs; i++){
        ans -= B[i] * (B[i] - 1) / 2;
    }

    cout << ans << endl;

    return 0;
}