【DSU編】AtCoder Library 解読 〜Pythonでの実装まで〜


# 0. はじめに
2020年9月7日に[AtCoder](https://atcoder.jp/)公式のアルゴリズム集 [AtCoder Library](https://atcoder.jp/posts/517) (**ACL**)が公開されました。
私はACLに収録されているアルゴリズム、データ構造のほとんどが初見だったのでいい機会だと思い、アルゴリズムの勉強からPythonでの実装までを行いました。

この記事では**DSU**をみていきます。

## 対象としている読者
* DSUってなに？という方。
* ACLのコードを見てみたけど**何をしているのかわからない**方。
* C++はわからないのでPythonで読み進めたい方。


## 参考にしたもの
AtCoder公式によるわかりやすい解説があります。

* [スライド](https://www.slideshare.net/chokudai/union-find-49066733/1)
* [youtube解説動画](https://www.youtube.com/watch?v=TdR816rqc3s&t=6822s)（ABC157の解説動画です）


# 1. DSUとは
**DSU** (**Disjoint Set Union**, **素集合データ構造**)はあるデータ集合を素集合(グループ)に分割して保持するデータ構造です。すなわち各データが１つのグループに属し、２つ以上のグループに属することはありません。

![dsu_1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/711140/8ef10d75-eeea-77f8-6946-7c115f7fd13f.png)

このデータ構造は以下の２つの便利な操作をサポートしています。

* **Union**: ２つのグループを１つに統合する。
* **Find**: ある要素がどのグループに属するかを求める。

このことから、このデータ構造を**UnionFind**と呼ぶこともあります。こちらの名称の方が馴染みがあるかもしれません。

![dsu_2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/711140/830d75f3-95f3-ad19-2a3a-01db833d78ee.png)

## 実装上は...
Unionはグループではなく要素を２つ指定し、それらの属するグループを統合します。また、各グループ内で１つの要素をリーダーとして選び、これをいわば”グループ名”として管理します。(上図での”グループ１”や”グループ２”は例えば１や４になります。)


# 2. 実装
それでは実装していきます。変数名、メソッド名等はなるべくACLに沿って実装します。

## 2.1. コンストラクタ
まずクラスDSUを作成し、コンストラクタを実装します。

```Python

class DSU:
    def __init__(self, n):  # n:要素数
        self._n = n
        self.parent_or_size = [-1] * n
```

nは要素数で、インスタンス変数 `_n`に保持しておきます。
また、各要素についての情報を格納するリスト`parent_or_size`を作成します。名前の示す通り、このリストの各要素`parent_or_size[i]`は、要素 i の parent (リーダー)もしくは要素 i の属するグループの size (大きさ)を表します。具体的には

<dl>
  <dt><code>parent_or_size[i]</code> が<strong>負</strong> : </dt>
  <dd> 要素 i はグループのリーダーであり、属するグループの大きさは <code>abs(parent_or_size[i])</code> </dd>
</dl>
<dl>
  <dt><code>parent_or_size[i]</code> が<strong> 0以上</strong> : </dt>
  <dd>要素 i の属するグループのリーダーは<code>parent_or_size[i]</code></dd>
</dl>

です。つまり初期化時点では全ての要素が、「自分をリーダーとする大きさ１のグループ(要素が自分だけのグループ)」に属しているということになります。

## 2.2. Find
操作FindはACLでは`leader`という名前で実装されています。このメソッドは要素を指定することで、その要素が属するグループのリーダーを返します。

```python
def leader(self, a):
    assert 0 <= a < self._n
    if self.parent_or_size[a] < 0: return a  # 負ならリーダー
    self.parent_or_size[a] = self.leader(self.parent_or_size[a])
    return self.parent_or_size[a]
```

まず、3行目で`parent_or_size[a]`が負ならば a はリーダーなのでそのままreturnします。そうでないならば`parent_or_size[a]`は a のリーダー **(仮)** です。なぜなら、グループの統合があった場合片方のリーダーはリーダーではなくなるからです。そこで、このリーダー(仮)から再帰的に **真の** リーダーを探します。4行目で右辺を直接returnせずに代入することでリーダーの情報を更新しています。メモ化していると言ってもいいかもしれません。これによって次回以降の探索が短くなります。

## 2.3. Union
操作UnionはACLでは`merge`という名前で実装されています。このメソッドは要素を２つ指定することで、それらの属するグループを統合します。

```python
def merge(self, a, b):
    assert 0 <= a < self._n
    assert 0 <= b < self._n
    x, y = self.leader(a), self.leader(b)
    if x == y: return x
    if -self.parent_or_size[x] < -self.parent_or_size[y]: x, y = y, x
    self.parent_or_size[x] += self.parent_or_size[y]  # xの大きさにyの大きさを加算
    self.parent_or_size[y] = x  # yのリーダーはx
    return x
```

まず、a, bそれぞれのリーダーx, yを求めます。これらが一致していれば、すでにa, bは同じグループに属しているので何もする必要はありません(5行目)。そうでない場合、yのグループをxのグループに統合します。この時、6行目のようにすることで常に小さいグループを大きいグループへ統合するようにします。yのグループメンバーはリーダーの情報を更新する必要があるので、これによって計算量を減らすことができます。

## 2.4. その他のメソッド
ACLのDSUには他にもいくつかのメソッドが実装されているので、それらも実装していきます。

### same
メソッド`same`は要素を２つ指定することで、それらが同じグループに属しているかの真偽値を返します。

```python
def same(self, a, b):
    assert 0 <= a < self._n
    assert 0 <= b < self._n
    return self.leader(a) == self.leader(b)
```

### size
メソッド`size`は指定した要素が属するグループの大きさ(要素数)を返します。

```python
def size(self, a):
    assert 0 <= a < self._n
    return -self.parent_or_size[self.leader(a)]
```

### groups
メソッド`groups`は全要素をグループごとにまとめたリストを返します。

```python
def groups(self):
    leader_buf = [self.leader(i) for i in range(self._n)]
    result = [[] for _ in range(self._n)]
    for i in range(self._n): result[leader_buf[i]].append(i)
    return [r for r in result if r != []]
```

## 2.5. まとめ

全てをまとめたものを置いておきます。

```python:dsu.py
class DSU:
    def __init__(self, n):
        self._n = n
        self.parent_or_size = [-1] * n
    
    def merge(self, a, b):
        assert 0 <= a < self._n
        assert 0 <= b < self._n
        x, y = self.leader(a), self.leader(b)
        if x == y: return x
        if -self.parent_or_size[x] < -self.parent_or_size[y]: x, y = y, x
        self.parent_or_size[x] += self.parent_or_size[y]
        self.parent_or_size[y] = x
        return x

    def same(self, a, b):
        assert 0 <= a < self._n
        assert 0 <= b < self._n
        return self.leader(a) == self.leader(b)

    def leader(self, a):
        assert 0 <= a < self._n
        if self.parent_or_size[a] < 0: return a
        self.parent_or_size[a] = self.leader(self.parent_or_size[a])
        return self.parent_or_size[a]
    
    def size(self, a):
        assert 0 <= a < self._n
        return -self.parent_or_size[self.leader(a)]
    
    def groups(self):
        leader_buf = [self.leader(i) for i in range(self._n)]
        result = [[] for _ in range(self._n)]
        for i in range(self._n): result[leader_buf[i]].append(i)
        return [r for r in result if r != []]
```



# 3. 使用例

```python
n = 8  # 全要素数
d = DSU(n)
d.merge(0, 1)
d.merge(1, 3)
d.merge(0, 4)
d.merge(5, 6)
d.merge(3, 7)
print(d.groups())  # [[0, 1, 3, 4, 7], [2], [5, 6]]
print(d.leader(3))  # 0
print(d.same(1, 7))  # True
print(d.same(0, 5))  # False
print(d.size(6))  # 2
```


# 4. 問題例
[AtCoder Library Practice Contest A "Disjoint Set Union"](https://atcoder.jp/contests/practice2/tasks/practice2_a)
[AtCoder Typical Contest 001 B "Union Find"](https://atcoder.jp/contests/atc001/tasks/unionfind_a)


# 5. おわりに
DSUの仕組みの解明からPythonでの実装までができました。また、実装上では様々な工夫がされていることがわかりました。特に必要な情報を一つの配列で表現する工夫には感動しました。

説明の間違いやバグ、アドバイス等ありましたらお知らせください。