公式

C - チームの旗の色 / Team Flag Colors 解説 by MtSaka


チームを合併していく操作を繰り返すため、Union Find(DSU)で情報を管理することを考えます。このとき、同じチームであることとUnion Find上で同一連結成分であることは同値です。

\(\text{col}[i]\) を 頂点 \(i\) が自身の連結成分の代表元の時に、その連結成分の色を指すとします。最初はすべての \(i\) について \(\text{col}[i]=-1\) として、色が設定されてないことを示します。

二つの連結成分を合併すると、新しく代表元ができます。その代表元を \(v\) とすると、\(\text{col}[v]\) を上書きされる色で更新します。

最後に、答えを求める部分では、すべての \(i\) について頂点 \(i\) が自身の連結成分の代表元であるか判定し、\(\text{col}[i]\)\(-1\) でない場合は色が設定されているため、それらを std::set などを用いて種類数を計算できます。

よって、時間計算量 \(\mathrm{O}(N\log N)\) などで解くことができます。 AtCoder Library のatcoder::dsu などを使うと実装しやすくなります。

実装例(C++)

#include <bits/stdc++.h>
#include <atcoder/dsu>
using namespace std;
int main() {
    int n, m;
    cin >> n >> m;
    vector<int> col(n, -1);
    atcoder::dsu uf(n);
    for (int i = 0; i < m; ++i) {
        int u, v, c;
        cin >> u >> v >> c;
        u--, v--;
        int nr = uf.merge(u, v);
        col[nr] = c;
    }
    set<int> st;
    for (int i = 0; i < n; i++)
        if (uf.leader(i) == i && col[i] != -1) st.emplace(col[i]);
    cout << st.size() << endl;
}

投稿日時:
最終更新: