蟻本 p85 食物連鎖(Union Find)を図を書きながら理解する

蟻本を進めていて、p85のUnion Findを使った例題が理解できず苦労してます。図を書いたら少しだけ理解が進んだので、メモしておきます。（勉強中のため、間違っている部分があったらごめんなさい！）

なお本記事はUnion Find自体の解説はなくて、Union Findをどう使うかの解説にフォーカスしています。

問題（ざっくり）

N匹の動物がいて、1〜Nまでの番号がついています
動物はA, B, Cの3種類です
動物AはBを食べ、動物BはCを食べ、動物CはAを食べます（食物連鎖）
2種類の情報が、K個連続して与えられるので、与えられる情報に矛盾がある個数を出力してください
情報 1 ：xとyは同じ動物です
情報 2 ：xはyを食べます

制約
1 <= N <= 50000
0 <= K <= 100000

解法

NやKの大きさからすると、全探索ではとうてい間に合わないです。
Union Findを使って、以下のようなグループを考えます。

（以下、本から抜粋）

各動物 iについて3つの要素 i-A, i-B, i-Cをつくり、3xN個の要素でUnion-findをつくります。定義は以下です。

i-xは「iが種類xである場合」を表す

各グループは、それらが全て同時に起こることがわかっていることを表す

私はこれを読んだ時、？？？？となって全然理解できませんでした。
これをもうちょっとかみ砕いて理解するために、前提を整理しながら図にしてみます。

情報1, 2から分かることを整理する

いったんUnion Findは忘れて、まず前提を整理しておきます。
与えられる情報に着目すると、

情報1：ある動物とある動物が同じ種類か？
情報2：ある動物とある動物の食物連鎖がどっち方向か？

という情報が分かります。逆にこの情報からは

ある動物がA,B,Cのどれにあたるのか？

という情報は分かりません。つまり、各動物に対して種類がA,B,C全てのパターンを考えておく必要があるということです。

例えば動物xと動物yに対して、図を書くとこんな感じです。

f:id:nakashiii:20211215234544j:plain

図の通り、動物xと動物yは種類A, B, Cのいずれかに存在しています。

情報1：xとyは同じ動物

次に情報1の「xとyは同じ動物」という情報について考えます。この時考えられるのは、

xもyも種類A
xもyも種類B
xもyも種類C

の3通りです。この3通りを図にします。ここで、解法にあった「同時に起こることを同じグループにする（= 頂点に対して辺を張る）」というルールを使って図を書きます。

f:id:nakashiii:20211215234539j:plain

ここで「xとyは同じ動物」という情報が矛盾する時を考えてみます。
それは「xとyは異なる動物」となります。（なんだか当たり前ですね、、）
これを図で示すと「異なる種類の動物に辺がはられている状態」となります。

f:id:nakashiii:20211215234535j:plain

これで、情報1に対してどのような時に情報が矛盾しているか分かるようになりました。

情報2：xはyを食べる

次に情報2ですが、「xはyを食べる」という情報をどのように表すかイメージしづらいです。
まずは、この時考えられるパターンを列挙します。
ただ、まずはさっきの図から矛盾が生じないようにしたいので、動物zを導入して「動物yは動物zを食べる」という情報が与えられたとします。

動物y = 種類Aの時、動物z = 種類B
動物y = 種類Bの時、動物z = 種類C
動物y = 種類Cの時、動物z = 種類A

これが考えられる全パターンです。これをさっきの図に書き足していきます。適用するルールは情報1と同じで「同時に起こることを同じグループにする（= 頂点に対して辺を張る）」です。

f:id:nakashiii:20211215234530j:plain

これで情報2も図で表現することができました。

そして、情報2が矛盾する時は以下の2通りです。

食物連鎖の方向が逆になっている（今回の場合、zはyを食べる）
同じ種類の動物になっている（今回の場合、yとzが同じ種類）

これを図示するとこんな感じです

f:id:nakashiii:20211215234526j:plain

これで情報1と情報2ともに矛盾をどうやって判定すればいいかが分かりました。

Union Findで表現する

ここまでで各情報に対しての矛盾の判定方法が分かったので、あとはこれをどうUnion Findで表現するかです。 N匹の動物に対して、それぞれ3種類ずつパターンを考える必要があるので、ノード数をN*3とします。ここで以下のようなルールで動物を表現します。

種類AのN匹：0 〜 N-1
種類BのN匹：N 〜 2N-1
種類CのN匹：2N 〜 3N-1

図にするとこんな感じです。

f:id:nakashiii:20211215234521j:plain

これでUnion Findで与えられる情報を表現することができます。

実装の流れ

以下のような流れで実装します

与えられた情報に対して矛盾があるかを判定してカウント（条件と矛盾するノードどうしが同じグループになっていないか？）
矛盾がない場合のみ、与えられた情報からUnionFindで併合
（上記を情報の個数だけ繰り返し）

計算量はO(K * log N)のはず、、

コード

~~すいません、スタミナ切れでコードまで載せられませんでした、、~~
→ 書いたので載せました！蟻本や他の方の解説記事とあわせてご参照ください

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

class UnionFind {
 public:
  // 親の番号を格納する。親だった場合は-1
  vector<int> Parent;
  UnionFind(int N) { Parent = vector<int>(N, -1); }

  // Aがどのグループに属しているか調べる
  int root(int A) {
    if (Parent[A] < 0) return A;
    return Parent[A] = root(Parent[A]);
  }

  // 自分のいるグループの頂点数を調べる
  int size(int A) {
    return -Parent[root(A)];  //親をとってきたい
  }

  // AとBをくっ付ける
  bool unite(int A, int B) {
    // AとBを直接つなぐのではなく、root(A)にroot(B)をくっつける
    A = root(A);
    B = root(B);
    if (A == B) {
      //すでにくっついてるからくっ付けない
      return false;
    }
    // 大きい方(A)に小さいほう(B)をくっ付ける
    // 大小が逆だったらひっくり返す
    if (size(A) < size(B)) {
      swap(A, B);
    }
    // Aのサイズを更新する
    Parent[A] += Parent[B];
    // Bの親をAに変更する
    Parent[B] = A;
    return true;
  }
};

int main() {
  int n, k;
  cin >> n >> k;

  UnionFind uni(3 * n);  // 正しい情報をグルーピング
  int ans = 0;           // 間違った情報をカウント

  for (int i = 0; i < k; i++) {
    int type, x, y;
    cin >> type >> x >> y;

    // 動物が1〜nの範囲内か判定
    if (x < 1 || x > n || y < 1 || y > n) {
      ans++;
      continue;
    }

    if (type == 1) {
      // 情報1が矛盾している場合（動物が同じ種類ではない）
      if (uni.root(x - 1) == uni.root(n + y - 1) ||
          uni.root(x - 1) == uni.root(2 * n + y - 1)) {
        ans++;
        continue;
      } else {
        uni.unite(x - 1, y - 1);                  // 種類A
        uni.unite(n + x - 1, n + y - 1);          // 種類B
        uni.unite(2 * n + x - 1, 2 * n + y - 1);  // 種類C
      }
    } else {
      // 情報2が矛盾している場合（食物連鎖が間違っている or
      // 動物が同じ種類）
      // if (uni.root(x) == uni.root(2 * n + y - 1) ||
      if (uni.root(y - 1) == uni.root(n + x - 1) ||
          uni.root(x - 1) == uni.root(y - 1)) {
        ans++;
        continue;
      } else {
        uni.unite(x - 1, n + y - 1);          // 種類AはBを食べる
        uni.unite(n + x - 1, 2 * n + y - 1);  // 種類BはCを食べる
        uni.unite(2 * n + x - 1, y - 1);      // 種類CはAを食べる
      }
    }
  }

  cout << ans << endl;
}

感想

図にするとUnion Findの「気持ち」が少し理解できました。あと、解説記事は初めてですが、自分の備忘録としても考えの整理としてすごくよかったです。
ただ、1時間くらいで書けるかなと思っていたら、3時間くらいかかってしまいました（汗）
どうしても理解が難しい問題に絞って、マイペースで記事を書ければいいかなと思います。

nakashiiiの自由帳

自由に書きます