🤖
ABC 183 | F - Confluence
問題
考えたこと
初見ではUnion-Findで解けそうな気がしますがひねってある。
入力例1は以下のようになる。
Union-Findを利用しつつ、Union-Findのグループのルートノードであるリーダーが所属するメンバーのクラスごとの生徒数を所持していればクエリ2を高速に処理できそう。
リーダのクラスごとのメンバ数はvector<vector<int>>
でv[i][j]
でリーダーunordered_map
を利用すればいい。
Union-Findで生徒が生徒が合流するとき以下のように処理し、すでに同じグループのときは何もしなければいい。
コード
実装時のTips
- 0-indexedで処理する
- Union-FindはAtCoderライブラリの
dsu
を利用する
#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;
int main() {
ll n, q;
cin >> n >> q;
vector<int> c(n);
vector<unordered_map<int, int>> nm(n);
for (int i = 0; i < n; i++) {
cin >> c[i];
c[i]--;
nm[i][c[i]]++;
}
dsu d(n);
for (int i = 0; i < q; i++) {
int t;
cin >> t;
if (t == 1) {
int a, b;
cin >> a >> b;
a--;
b--;
if (d.same(a, b)) {
continue;
}
int al = d.leader(a);
int bl = d.leader(b);
int to = d.merge(a, b);
int from = (al + bl) - to;
for (auto [cls, cnt] : nm[from]) {
nm[to][cls] += cnt;
}
} else {
int x, y;
cin >> x >> y;
x--;
y--;
ll v = d.leader(x);
cout << nm[v][y] << endl;
}
}
}
Discussion