AtCoder Library (の一部)を Cython で ラップ して Python で使えるようにする

# はじめに
この記事は、[AtCoder Library](https://atcoder.jp/posts/517) のうち、[DSU](https://atcoder.github.io/ac-library/production/document_ja/dsu.html) と [Fenwick Tree](https://atcoder.github.io/ac-library/production/document_ja/fenwicktree.html) を、Cython を利用して Python から使えるようにする記事となります。
DSU と Fenwick Tree 以外で、同じ方法で利用できるものがあるかは確認していませんが、例えば、[Segtree](https://atcoder.github.io/ac-library/production/document_ja/segtree.html) は非型テンプレートを使っていて、Cython は非型テンプレートに対応していない (多分) ので難しいと思います。

AtCoder Library は C++ で書かれているので、Python で 同じ処理を実装をしたときよりも速くなります。
しかし、後述しますが、PyPy と比較した場合は、速いとは言えない結果になりました。


なお、本記事では Cython の文法などについてはほとんど触れませんので、この記事で Cython について気になった方は他の方の記事も参考にしてみてください。(Cython は、Python のコードをそのまま書いても一応動きます。)
# Cython とは
C/C++ を利用して Python の拡張モジュールを作るための言語で、C/C++ のクラスや関数などを呼び出したりすることができます。
まさに今回のようなことをするための言語です。(本当ですか？)
# 環境構築
Cython と AtCoder Library が必要なので、インストールなどをしていきます。
なお、Windows の人は、Cython のコンパイルで躓くことがあるので、できれば Linux や WSL を使うことをおすすめします。

## Cython のインストール

```sh:terminal
$ pip3 install cython
```

で行います。
Cython が使えるか確かめてみましょう。


```cython:hello.pyx
print("Hello, World!")
```

Hello, World! を出力するプログラムです。拡張子が`.pyx`であることに注意してください。
これをコンパイルして実行してみます。

```sh:terminal
$ cythonize -i -3 hello.pyx
$ python3 -c "import hello"
```

1 行目の`cythonize`コマンドでコンパイルしています。
`-i`オプションは、まず C のコードに変換して、それを Python で import できる形式 (Linux なら`.so`、Windows なら`.pyd` ) にコンパイルします。
`-3`オプションは Python が 3 系のときに使います。

2 行目のように、Python で hello を import するコードを実行し、 Hello, World! が出力されれば、Cython の環境構築は成功です。
## AtCoder Library
`/opt/ac-library/` フォルダを作ります。(AtCoder のジャッジ環境に合わせています。)
その直下に https://img.atcoder.jp/practice2/ac-library.zip の中の、`atcoder`フォルダをコピーします。

# AtCoder Library を Cython で ラップ する
## DSU
まずは [DSU](https://atcoder.github.io/ac-library/production/document_ja/dsu.html) を使えるようにします。
Python から AtCoder Library を利用するためには、Cython で次の 2 ステップが必要です。
1. Cython から AtCoder Library を使えるようにする
2. Python からも使えるようにラップする
### 1. Cython から AtCoder Library を使えるようにする

```cython
# distutils: language=c++
# distutils: include_dirs=/opt/ac-library
from libcpp cimport bool
from libcpp.vector cimport vector
cdef extern from "<atcoder/dsu>" namespace "atcoder":
    cdef cppclass dsu:
        dsu(int n) # コンストラクタ
        int merge(int a, int b) # 辺 (a, b) を足す
        bool same(int a, int b) # 頂点 a, b が連結かどうか
        int leader(int a) # 頂点 a の属する連結成分の代表元
        int size(int a) # 頂点 a の属する連結成分のサイズ
        vector[vector[int]] groups() # 「一つの連結成分の頂点番号のリスト」のリスト
```
これは DSU を Cython で使えるようにしたコードです。
1 行目は、C++ を利用することを、2 行目は、AtCoder Library の場所をそれぞれ指定しています。
3、4 行目は C++ の bool と vector を利用するために import しています。
5 行目以降は、`atcoder/dsu.hpp`の`atcoder`名前空間にある`dsu`クラスのメソッド (正直 C++ はよくわからないので、この辺は各自で内部実装や、[DSU のドキュメント](https://atcoder.github.io/ac-library/production/document_ja/dsu.html)を読んでみてください……。) のうち、Cython で使うものを書いています。
これで、Cython 内では DSU を

```cython
cdef dsu *u = new dsu(3)
u.merge(0, 1)
print(u.same(0, 1)) # True
print(u.same(0, 2)) # False
```

のように利用できるようなりました。
### 2. Python からも使えるようにラップする

```cython
cdef class DSU:
    cdef dsu *_thisptr
    # コンストラクタ
    def __cinit__(self, int n):
        self._thisptr = new dsu(n)
    # 辺 (a, b) を足す
    cpdef int merge(self, int a, int b):
        return self._thisptr.merge(a, b)
    # 頂点 a, b が連結かどうか
    cpdef bool same(self, int a, int b):
        return self._thisptr.same(a, b)
    # 頂点 a の属する連結成分の代表元
    cpdef int leader(self, int a):
        return self._thisptr.leader(a)
    # 頂点 a の属する連結成分のサイズ
    cpdef int size(self, int a):
        return self._thisptr.size(a)
    # 「一つの連結成分の頂点番号のリスト」のリスト
    cpdef vector[vector[int]] groups(self):
        return self._thisptr.groups()
```

cpdef で定義された 関数やメソッドは Python からも呼び出すことができるようになります。
DSU クラスを作り、内部では dsu クラスのメソッドを呼び出すことで、Python からの DSU の利用を実現しているということになります。

これで Cython のコードは完成です。

```cython:dsu.pyx
# distutils: language=c++
# distutils: include_dirs=/opt/ac-library
from libcpp cimport bool
from libcpp.vector cimport vector
cdef extern from "<atcoder/dsu>" namespace "atcoder":
    cdef cppclass dsu:
        dsu(int n)
        int merge(int a, int b)
        bool same(int a, int b)
        int leader(int a)
        int size(int a)
        vector[vector[int]] groups()
cdef class DSU:
    cdef dsu *_thisptr
    def __cinit__(self, int n):
        self._thisptr = new dsu(n)
    cpdef int merge(self, int a, int b):
        return self._thisptr.merge(a, b)
    cpdef bool same(self, int a, int b):
        return self._thisptr.same(a, b)
    cpdef int leader(self, int a):
        return self._thisptr.leader(a)
    cpdef int size(self, int a):
        return self._thisptr.size(a)
    cpdef vector[vector[int]] groups(self):
        return self._thisptr.groups()
```

Python で import できるようにコンパイルしましょう。

```sh:terminal
$ cythonize -i -3 dsu.pyx
```

### Python で実行
Python で [DSU の練習問題](https://atcoder.jp/contests/practice2/tasks/practice2_a)を解くコードを書いてみます。

```python3:ALPC-A.py
from dsu import DSU
n, q = map(int, input().split())
dsu = DSU(n)
res = []
for i in range(q):
    t, u, v = map(int, input().split())
    if t == 0:
        dsu.merge(u, v)
    else:
        if dsu.same(u, v):
            res.append(1)
        else:
            res.append(0)
print('\n'.join(map(str, res)))
```

練習問題のサンプルを入力してみて、正しい出力を得られたら、とりあえず成功です。
### 提出
実際に [DSU の練習問題](https://atcoder.jp/contests/practice2/tasks/practice2_a)に提出してみて、AC になることを確認したいのですが、AtCoder のジャッジ側に`dsu.pyx`やそれをコンパイルしたモジュールがあるわけではないので、そのまま提出しても RE になってします。
そこで、提出コードに`dsu.pyx`を埋め込む必要があります。

```python3:ALPC-A.py
cycode = r"""
# distutils: language=c++
# distutils: include_dirs=/opt/ac-library
from libcpp cimport bool
from libcpp.vector cimport vector
cdef extern from "<atcoder/dsu>" namespace "atcoder":
    cdef cppclass dsu:
        dsu(int n)
        int merge(int a, int b)
        bool same(int a, int b)
        int leader(int a)
        int size(int a)
        vector[vector[int]] groups()
cdef class DSU:
    cdef dsu *_thisptr
    def __cinit__(self, int n):
        self._thisptr = new dsu(n)
    cpdef int merge(self, int a, int b):
        return self._thisptr.merge(a, b)
    cpdef bool same(self, int a, int b):
        return self._thisptr.same(a, b)
    cpdef int leader(self, int a):
        return self._thisptr.leader(a)
    cpdef int size(self, int a):
        return self._thisptr.size(a)
    cpdef vector[vector[int]] groups(self):
        return self._thisptr.groups()
"""
import os, sys
if sys.argv[-1] == 'ONLINE_JUDGE':
    open('dsu.pyx', 'w').write(cycode)
    os.system('cythonize -i -3 dsu.pyx')
from dsu import DSU
n, q = map(int, input().split())
dsu = DSU(n)
res = []
for i in range(q):
    t, u, v = map(int, input().split())
    if t == 0:
        dsu.merge(u, v)
    else:
        if dsu.same(u, v):
            res.append(1)
        else:
            res.append(0)
print('\n'.join(map(str, res)))
```

AtCoder の Python は、Numba の AOT (詳しくは[こちらの記事](https://qiita.com/yniji/items/d012bb9f938e0445a3ff#numba-aot)を読んでください) のためにコンパイルフェーズ (このときかかる時間は実行時間に含まれない) が用意されていて、コンパイルフェーズ時のコマンドライン引数には`ONLINE_JUDGE`が与えられます。 
これを利用し、30 行目の`if sys.argv[-1] == 'ONLINE_JUDGE':`で、コンパイルフェーズかどうかを判断し、コンパイルフェーズだった場合、`dsu.pyx`のファイル書き込みと、`cythonize`コマンドによるコンパイルを行います。
このコードを[提出](https://atcoder.jp/contests/practice2/submissions/17535199)して、424 ms で無事 AC になりました。
比較として、有志による、AtCoder Library の [Python 移植版](https://github.com/not522/ac-library-python)の [DSU](https://github.com/not522/ac-library-python/blob/master/atcoder/dsu.py) を使った場合、[Python で提出](https://atcoder.jp/contests/practice2/submissions/17535364)したときは、635 ms、[PyPy で提出](https://atcoder.jp/contests/practice2/submissions/17535461)したときは 551 ms でした。
~~あれ、苦労したわりにそんなに速くなってねぇんだが~~
## Fenwick Tree
同様に [Fenwick Tree](https://atcoder.github.io/ac-library/production/document_ja/fenwicktree.html) も Python で使えるようにします。

```cython:fenwick_tree.pyx
# distutils: language=c++
# distutils: include_dirs=/opt/ac-library
cdef extern from "<atcoder/fenwicktree>" namespace "atcoder":
    cdef cppclass fenwick_tree[T]:
        fenwick_tree(int n)
        void add(int p, T x)
        T sum(int l, int r)
ctypedef long long ll
cdef class FenwickTree:
    cdef fenwick_tree[ll] *_thisptr
    def __cinit__(self, int n):
        self._thisptr = new fenwick_tree[ll](n)
    cpdef void add(self, int p, ll x):
        self._thisptr.add(p, x)
    cpdef ll sum(self, int l, int r):
        return self._thisptr.sum(l, r)
```
基本的には DSU と同じですが、Fenwick Tree はテンプレートが使われています。
[ドキュメント](https://atcoder.github.io/ac-library/production/document_ja/fenwicktree.html)によると、`int / uint (unsigned int) / ll (long long) / ull (unsigned long long) / modint`を型に指定できるようです。
ここでは、型を ll にしています。(int と uint は ll で代用できると思うので)
ull の Fenwick Tree が欲しかったら、書き換えるか、もう 1 つクラスを作ってください。
~~modint は諦めました……。~~

[Fenwick Tree の練習問題](https://atcoder.jp/contests/practice2/tasks/practice2_b)を解いてみます。

```python3:ALPC-B.py
cycode = r"""
# distutils: language=c++
# distutils: include_dirs=/opt/ac-library
cdef extern from "<atcoder/fenwicktree>" namespace "atcoder":
    cdef cppclass fenwick_tree[T]:
        fenwick_tree(int n)
        void add(int p, T x)
        T sum(int l, int r)
ctypedef long long ll
cdef class FenwickTree:
    cdef fenwick_tree[ll] *_thisptr
    def __cinit__(self, int n):
        self._thisptr = new fenwick_tree[ll](n)
    cpdef void add(self, int p, ll x):
        self._thisptr.add(p, x)
    cpdef ll sum(self, int l, int r):
        return self._thisptr.sum(l, r)
"""
import os, sys
if sys.argv[-1] == 'ONLINE_JUDGE':
    open('fenwick_tree.pyx', 'w').write(cycode)
    os.system('cythonize -i -3 fenwick_tree.pyx')
from fenwick_tree import FenwickTree
n, q = map(int, input().split())
ft = FenwickTree(n)
for i, a in enumerate(map(int, input().split())):
    ft.add(i, a)
res = []
for _ in range(q):
    t, u, v = map(int, input().split())
    if t == 0:
        ft.add(u, v)
    else:
        res.append(ft.sum(u, v))
print('\n'.join(map(str, res)))
```

このコードを[提出](https://atcoder.jp/contests/practice2/submissions/17535413)して、1287 ms で AC になりました。
比較として、[Python 移植版の Fenwick Tree](https://github.com/not522/ac-library-python/blob/master/atcoder/fenwicktree.py) を使った場合、[Python で提出](https://atcoder.jp/contests/practice2/submissions/17535420)したときは、4663 ms、[PyPy で提出](https://atcoder.jp/contests/practice2/submissions/17535464)したときは 1001 ms でした。
~~PyPyより遅いやんけ！キレてしまった……。~~
# 比較
## DSU
| Cython + Python | Python | PyPy |
|:-:|:-:|:-:|
| **424 ms** | 635 ms | 551 ms |
## Fenwick Tree
| Cython + Python | Python | PyPy |
|:-:|:-:|:-:|
| 1287 ms | 4663 ms | **1001 ms** |
# おまけ
実は、Python で行っていた部分の処理も Cython で書いた場合は、PyPy よりもずっと速くなります。
以下は [Fenwick Tree の練習問題](https://atcoder.jp/contests/practice2/tasks/practice2_b)を 全部 Cython で書いた場合のコードです。 

```cython:ALPC-B.pyx
# distutils: language=c++
# distutils: include_dirs=/opt/ac-library
from libc.stdio cimport scanf, printf
from libcpp.vector cimport vector
cdef extern from "<atcoder/fenwicktree>" namespace "atcoder":
    cdef cppclass fenwick_tree[T]:
        fenwick_tree(int n)
        void add(int p, T x)
        T sum(int l, int r)
ctypedef long long ll
cdef int n, q, i, t
cdef ll u, v, a
scanf('%d%d', &n, &q)
cdef fenwick_tree[ll] *ft = new fenwick_tree[ll](n)
for i in range(n):
    scanf('%lld', &a)
    ft.add(i, a)
for i in range(q):
    scanf('%d%lld%lld', &t, &u, &v)
    if t == 0:
        ft.add(u, v)
    else:
        printf('%lld\n', ft.sum(u, v))
```

この[提出](https://atcoder.jp/contests/practice2/submissions/17536517)は 251 ms でした。

| Cython | Cython + Python | Python | PyPy |
|:-:|:-:|:-:|:-:|
| **251 ms** | 1287 ms | 4663 ms | 1001 ms |
超速い。
ではなぜ Cython + Python は PyPy より遅くなったのでしょうか。
原因は Python のループ処理の遅さにあるのでは無いかと考えました。
実験として、この問題の入力を受け取るだけのコードを Python と PyPy それぞれで提出して実行時間を見てみます。

```python3:ALPC-B-input.py
n, q = map(int, input().split())
for i, a in enumerate(map(int, input().split())):
    i, a
for _ in range(q):
    t, u, v = map(int, input().split())
```

結果は [Python](https://atcoder.jp/contests/practice2/submissions/17536913) は 1012 ms、[PyPy](https://atcoder.jp/contests/practice2/submissions/17536915) は 687 ms でした。
AC コードの実行時間からこの結果を引くと、Fenwick Tree の処理の部分の時間は大体 Cython + Python は 275ms、PyPy は 314ms くらいかかっていそうです。~~微妙~~

# まとめ
PyPy 使えばいいと思う。