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

はじめに

この記事は、AtCoder Library のうち、DSU と Fenwick Tree を、Cython を利用して Python から使えるようにする記事となります。
DSU と Fenwick Tree 以外で、同じ方法で利用できるものがあるかは確認していませんが、例えば、Segtree は非型テンプレートを使っていて、Cython は非型テンプレートに対応していない (多分) ので難しいと思います。

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

なお、本記事では Cython の文法などについてはほとんど触れませんので、この記事で Cython について気になった方は他の方の記事も参考にしてみてください。(Cython は、Python のコードをそのまま書いても一応動きます。)

Cython とは

C/C++ を利用して Python の拡張モジュールを作るための言語で、C/C++ のクラスや関数などを呼び出したりすることができます。
まさに今回のようなことをするための言語です。(本当ですか？)

環境構築

Cython と AtCoder Library が必要なので、インストールなどをしていきます。
なお、Windows の人は、Cython のコンパイルで躓くことがあるので、できれば Linux や WSL を使うことをおすすめします。

Cython のインストール

terminal

$ pip3 install cython

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

hello.pyx

print("Hello, World!")

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

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 を使えるようにします。
Python から AtCoder Library を利用するためには、Cython で次の 2 ステップが必要です。

1. Cython から AtCoder Library を使えるようにする
2. Python からも使えるようにラップする

1. Cython から AtCoder Library を使えるようにする

# 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 のドキュメントを読んでみてください……。) のうち、Cython で使うものを書いています。
これで、Cython 内では DSU を

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

のように利用できるようなりました。

2. Python からも使えるようにラップする

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 のコードは完成です。

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 できるようにコンパイルしましょう。

terminal

$ cythonize -i -3 dsu.pyx

Python で実行

Python で DSU の練習問題を解くコードを書いてみます。

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 の練習問題に提出してみて、AC になることを確認したいのですが、AtCoder のジャッジ側にdsu.pyxやそれをコンパイルしたモジュールがあるわけではないので、そのまま提出しても RE になってします。
そこで、提出コードにdsu.pyxを埋め込む必要があります。

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 (詳しくはこちらの記事を読んでください) のためにコンパイルフェーズ (このときかかる時間は実行時間に含まれない) が用意されていて、コンパイルフェーズ時のコマンドライン引数にはONLINE_JUDGEが与えられます。
これを利用し、30 行目のif sys.argv[-1] == 'ONLINE_JUDGE':で、コンパイルフェーズかどうかを判断し、コンパイルフェーズだった場合、dsu.pyxのファイル書き込みと、cythonizeコマンドによるコンパイルを行います。
このコードを提出して、424 ms で無事 AC になりました。
比較として、有志による、AtCoder Library の Python 移植版の DSU を使った場合、Python で提出したときは、635 ms、PyPy で提出したときは 551 ms でした。
あれ、苦労したわりにそんなに速くなってねぇんだが

Fenwick Tree

同様に Fenwick Tree も Python で使えるようにします。

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 はテンプレートが使われています。
ドキュメントによると、int / uint (unsigned int) / ll (long long) / ull (unsigned long long) / modintを型に指定できるようです。
ここでは、型を ll にしています。(int と uint は ll で代用できると思うので)
ull の Fenwick Tree が欲しかったら、書き換えるか、もう 1 つクラスを作ってください。
modint は諦めました……。

Fenwick Tree の練習問題を解いてみます。

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)))

このコードを提出して、1287 ms で AC になりました。
比較として、Python 移植版の Fenwick Tree を使った場合、Python で提出したときは、4663 ms、PyPy で提出したときは 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 の練習問題を 全部 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))

この提出は 251 ms でした。

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

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 は 1012 ms、PyPy は 687 ms でした。
AC コードの実行時間からこの結果を引くと、Fenwick Tree の処理の部分の時間は大体 Cython + Python は 275ms、PyPy は 314ms くらいかかっていそうです。微妙

まとめ

PyPy 使えばいいと思う。