Zenn
🤖

ABC 183 | F - Confluence

2020/11/16に公開
#
map
#
競プロ
#
unionfind
#
辞書
tech

問題

https://atcoder.jp/contests/abc183/tasks/abc183_f

考えたこと

初見ではUnion-Findで解けそうな気がしますがひねってある。

入力例1は以下のようになる。

Union-Findを利用しつつ、Union-Findのグループのルートノードであるリーダーが所属するメンバーのクラスごとの生徒数を所持していればクエリ2を高速に処理できそう。
リーダのクラスごとのメンバ数はvector<vector<int>>でv[i][j]でリーダーiのグループのクラスjの人数を管理すると巨大になりメモリオーバーでREになってしまう。このようなスパースなデータの場合は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