【Cython】C++並みに速い競プロのPython高速化

# はじめに

競プロではPythonは遅くていじられがちですが、高速化手段が豊富なため決して不利ではありません。

PyPyより速い高速化のうち、Cythonは最も汎用性が高く競プロ向きだと思いますが、文献が非常に少なく普及していないという弱点があります。

そのため、この記事では、Cythonの高速化方法をまとめ、Cythonを広めることを目標にしています。

## Cythonとは
Pythonと親和性の高い構文で、PythonのコードにC/C++のコードを混ぜ込むことができるPython拡張です。

C/C++が含まれるので、事前にコンパイルが必要があります。

Pythonに機能を追加した感じなので、基本的にPythonのコードはそのまま使用できます。

### Cythonの威力

https://rheotommy.hatenablog.com/entry/2020/07/18/205343

をCythonで高速化してみます。

※AtCoderでCythonを使うときのフォーマットを使用しているので、本質部分は`mycode`のみです。

```py
import sys
if sys.argv[-1] == 'ONLINE_JUDGE':
  mycode = r'''
# cython: language_level = 3, nonecheck = False, cdivision = True
ctypedef unsigned long long ull

cdef ull collatz(ull i) nogil:
  cdef ull cnt = 0
  while i != 1:
    cnt += 1
    if i % 2 == 0:
      i //= 2
    else:
      i *= 3
      i += 1
  return cnt

def main():
  cdef ull i, n = int(input()), acc = 0
  with nogil:
    for i in range(1, n + 1):
      acc += collatz(i)
      acc %= 1000000007
  print(acc)

if __name__ == 'mycode':
  main()
'''
  with open('mycode.pyx', 'w') as f:
    f.write(mycode)
  setup = r'''
from distutils.core import setup, Extension
from Cython.Build import cythonize
from numpy import get_include
ext = Extension(
  'mycode',
  sources = ['mycode.pyx'],
  include_dirs = ['.', get_include()],
  extra_compile_args = ['-O3']
)
setup(name = 'mycode', ext_modules = cythonize([ext]))
'''
  with open('setup.py', 'w') as f:
    f.write(setup)
  import subprocess
  with open('error.txt', 'w') as ef:
    code = subprocess.run('python3.8 setup.py build_ext --inplace'.split(), stderr = ef).returncode
  if code:
    with open('error.txt', 'r') as ef:
      with open('mycode.py', 'w') as f:
        f.write(f'''
import sys
sys.stderr.write("""{ef.read()}""")
exit({code})
''')
  exit()
import mycode
```

$10 ^ 7$ で2200msだったので、普通に速い感じです。

# AtCoderでCythonを使うときのフォーマット

※**2023年のAtCoder言語アップデートにより、基本的にこのフォーマットは不要になりました。**

AtCoderでは、CythonはCにコンパイルされてしまいます。

すると、vectorやunordered_mapといった競プロで必須のC++のデータ構造を使うことができなくなり、Cythonは使い物にならなくなります。

そこで、以下のコードを**Python**で提出します：

```py
import sys
if sys.argv[-1] == 'ONLINE_JUDGE':
  mycode = r'''
# distutils: language = c++
# cython: language_level = 3
def main():
  pass # ここに書く
if __name__ == 'mycode':
  main()
'''
  with open('mycode.pyx', 'w') as f:
    f.write(mycode)
  import subprocess
  with open('error.txt', 'w') as ef:
    code = subprocess.run('cythonize -i -3 -b mycode.pyx'.split(), stderr = ef).returncode
  if code:
    with open('error.txt', 'r') as ef:
      with open('mycode.py', 'w') as f:
        f.write(f'''
import sys
sys.stderr.write("""{ef.read()}""")
exit({code})
''')
  exit()
import mycode
```

※ [AtCoderでCythonの力を解放する魔術詠唱](https://aotamasaki.hatenablog.com/entry/2020/09/09/Cython_in_AtCoder) を参考にしています。

AtCoderのPythonは、他の言語と違い、事前に一度実行されるという特徴があります。そのため、事前に実行された時(ONLINE_JUDGEという引数が入る)はファイルを作成してコンパイルし、それ以降の実行では事前にコンパイルされた`mycode`を実行するようにしています。

また、コンパイルに失敗した場合、`mycode`を標準エラー出力にエラーメッセージを出力するコードに変えています。

`def main()` の中にいれると少し速いです。

# 静的型付け

静的型付けとは、コンパイル時に型が何か分かるように指定することです。Cythonにおいては、高速化方法のほとんどが静的型付けです。

### 基本構文

```pyx
cdef [型] [変数名1], [変数名2], ...
```

で静的型付けできます。構文はC/C++と同じです。また、次のようにインデントを使って書くこともできます。

```pyx
cdef:
  [型] [変数名]
  [型] [変数名]
  ...
```

ただし、for文、if文、while文の中で静的型付けはできません。Pythonにスコープがないからです。

### 静的型付けできるもの

- C/C++の型
- Pythonの組み込みの型
- `cdef class`で定義した自作の型

です。このうちPythonの組み込みデータ構造を型付けしてもあまり効果がないため、C/C++の型が中心になります。

### ctypedef

```pyx
ctypedef [元の型名] [型の別名]
```

で型名を別名で書けるようになります。

### C/C++の型についての知識

数値型については以下のサイトなどに説明されています。

http://www.kusa.ac.jp/~kajiura/c/data/newpage1.htm

また、C++のSTLについては、以下のサイトを参考にすると良いと思います。

http://vivi.dyndns.org/tech/cpp/container.html

## Cの型

Cの型は組み込みで用意されています。生配列や`const`、ポインタもあります。

### 型変換

Cの型とPythonの型には暗黙の型変換が定義されているので、Pythonの関数にCの型を渡すことができます(ただし変換を行うロスが発生します)。

### 注意点

- `bool`型は`bint`と書く
- `char`型は`int`との型変換があり、文字は`b'a'`などとして`bytes`型で書く必要がある
- 少し演算がC++と異なるところがある、`cython.operator`に全て用意されているのでわからなかったらそれを使う
  - `*ptr` → `ptr[0]`、`cython.operator.dereference(ptr)`
  - `ptr->func()` → `ptr.func()`
  - `ptr++` → `cython.operator.increment(ptr)`

### 高速化例

$10^9!$ を $10^9 + 7$ で割った余りを求めるプログラムで試してみます。

```py
ans = 1
for i in range(1, 1000000001): # 1から10 ^ 9
  ans = ans * i % 1000000007
print(ans)
```

このコードでPyPyだと5200msほどでした。

Cythonでは次のようなコードとなります。

```pyx
# distutils: language = c++
# cython: language_level = 3
ctypedef unsigned long long ull
def main():
  cdef ull ans = 1, i
  for i in range(1, 1000000001): # 1から10 ^ 9
    ans = ans * i % 1000000007
  print(ans)
if __name__ == 'mycode':
  main()
```

3300msほどで爆速です。なお、`range`は特殊で、型指定をするとPythonの関数ではなくなるので、遅くはありません。

## C++の型

### cimport

C++のデータ構造はSTLライブラリにあるので、`cimport`キーワードでimportする必要があります。`import`キーワードと基本的に使い方は同じです。

- Cのライブラリは`libc`内にあり、以下が使用可能です。

https://github.com/cython/cython/tree/master/Cython/Includes/libc

- C++のライブラリは`libcpp`にあり、以下が使用可能です。

https://github.com/cython/cython/tree/master/Cython/Includes/libcpp

例えば、`printf`と`vector`をimportするなら、

```pyx
from libc.stdio cimport printf
from libcpp.vector cimport vector
```

です。

### 型変換

こちらも暗黙の型変換が定義されていて、`map`と`dict`、`string`と`bytes`などが勝手に型変換されます。また、`vector`や`pair`といった型はPythonのiterableな型(`list`や`tuple`)ならどれとも型変換できます。

さらに、大体の型はPythonの組み込み型と似たような`str`型への変換があるので、printデバッグが容易にできます。

### 注意点

・名前がPythonの組み込み関数などと被ってしまうSTLは次のようにimportするといいです。

```pyx
from libcpp.map cimport map as cmap
```

・テンプレートは`<>`ではなく`[]`で書きます。

```pyx
cdef vector[int] a
```

・`string`型は`bytes`型と対応しています。

```pyx
cdef string s = b'Hello world!'
```

### 高速化例

[Educational DP Contest O - Matching](https://atcoder.jp/contests/dp/tasks/dp_o) を例にしてみます。

PyPyでは1300msほどです。

https://atcoder.jp/contests/dp/submissions/20684660

Cythonでは700msでした。

https://atcoder.jp/contests/dp/submissions/24437628

注目すべきは、型指定を最初に書いたことと、`dp[-1]`を`dp.back()`にしたこと以外変更がないことです。

入出力の高速化や、Pythonのリストの使用をやめるなど、更なる高速化はありますが、それらをせずにここまで高速化できます。
・for文内で回す変数
・配列などデータ構造
は型付けの効果が大きいです。

## 型推論

Cythonには`auto`はありませんが、**ローカルスコープでは**`cdef`を使用せずに**C/C++の関数の返り値**を代入しても、型推論され高速になります。

```pyx
# distutils: language = c++
# cython: language_level = 3
from libcpp.vector cimport vector
def main():
  cdef vector[int] a
  i = a.begin() # 型推論
  i = 0 # エラー
```

# 関数

### 基本構文

```pyx
cdef [返り値の型] func([引数の型] arg1, [引数の型] arg2, ...):
  ...
```

のように書きます。型名は省略可能です。参照渡しもあります。

### fused type

```pyx
ctypedef fused [fused type名]:
  [型・fused type]
  [型・fused type]
  ...
```

Cythonにはテンプレートがないので、ジェネリクスを実現するためにfused typeが存在します。型をまとめて扱えるようになります。

chmax関数を例に取ります。

```pyx
# distutils: language = c++
# cython: language_level = 3
ctypedef fused integer: # 整数
  bint
  char
  unsigned char
  short
  unsigned short
  int
  unsigned int
  long long
  unsigned long long
ctypedef fused real: # 実数
  integer
  float
  double
  long double
cdef bint chmax(real& a, real b):
  cdef real* ptr = &a # aを変更するには一度ポインタを取得する
  if a < b:
    ptr[0] = b
    return 1
  return 0
```

# class

### 基本構文

```pyx
cdef class Cls:
  ...
```

のようにすることでメンバ関数を`cdef`で定義することができます。また、クラス変数を`cdef`で静的型付けできます。

`Hello`クラスを例に取ります。

```pyx
# distutils: language = c++
# cython: language_level = 3
from libcpp.string cimport string
cdef class Hello:
  cdef string s
  def __init__(self, string s): # 特殊メソッドはdef、defも引数は型付け可能
    self.s = s
  cdef void say(self):
    print(self.s)
cdef Hello obj = Hello(b'Hello world!') # cdef型付け可能
```

# その他

### 関数で返される参照に代入する

例えば、イテレータ `itr` の参照先に代入したいときは、`cython.operator.dereference(itr) = 0` というようなことはPythonの構文上できません。そのため、ポインタを使うしか方法がないです。

```pyx
# distutils: language = c++
# cython: language_level = 3
from libcpp.vector cimport vector
cimport cython
def main():
  cdef vector[int] a = [1, 2, 3]
  itr = a.begin()
  cdef int* ptr = &cython.operator.dereference(itr)
  ptr[0] = 0 # ポインタの参照先になら代入できる
  print(a) # [0, 2, 3]
if __name__ == 'mycode':
  main()
```

### DEF・IF

`DEF`、`IF`というキーワードが用意されていて、`DEF`はコンパイル時に置き換えてくれるマクロで、`IF`はコンパイル時ifです。

```pyx
DEF inf = 2305843009213693952LL
```

### nogil

global interpreter lock(GIL)を切ってくれます。`with nogil:`とするか、あるいは`cdef void func() nogil:`のように`cdef`関数につけることができます。内部ではGILのある関数やクラスは一切使えなくなります。

### cpdef

`cdef`をつけるとCythonからは呼び出せますが、Pythonからは呼び出せなくなるので、どちらからも呼び出したい場合は`cpdef`を使用します。

### cdef extern from
実は`cimport`の内部は全てこれで実装されています。C/C++のライブラリを呼び出すことができます。Cythonの中心的な機能の一つです。

```pyx
cdef extern from '[ファイル名]' namespace '[名前空間]': # namespaceは省略可
  ...
```

と書き、インデントブロックのなかに、C++のプロトタイプ宣言のように関数、マクロ、クラスを書きます。また、`nogil`をつけることもできます。

クラスは`cppclass`と書き、テンプレートは`void func[T]`や`cppclass Cls[T]`のように書きます。それ以外はC++と同じ書き方です。

また、次のようにすることで別名をつけることもできます。

```pyx
cdef extern from '<algorithm>' nogil:
  void cppsort 'std::sort'[T](T, T) # std::sortをcppsortという名前で呼び出せるようになる
```

次はそれぞれ高速`input`関数、`iostream`の例です。

```pyx
# distutils: language = c++
# cython: language_level = 3
from libcpp.string cimport string
from libc.stdio cimport EOF
cdef extern from '<stdio.h>' nogil: # 名前空間なし
  int getchar_unlocked()
cdef str input():
  cdef string s
  cdef char c = getchar_unlocked()
  while c != b'\n' and c != EOF:
    s += c
    c = getchar_unlocked()
  return s.decode() # bytes -> str
```

```pyx
# distutils: language = c++
# cython: language_level = 3
cdef extern from '<iostream>' namespace 'std' nogil:
  cppclass ostream:
    ostream& operator <<[T](T)
  ostream& operator <<[T](ostream&, const T&)
  ostream cout, endl
  cppclass istream:
    istream& operator >>[T](T&)
    ostream* tie(ostream*)
  istream cin
  # static関数は、@staticmethodを使うことができる
  cppclass ios_base:
    @staticmethod
    bint sync_with_stdio(bint)
cin.tie(NULL) # cin高速化
ios_base.sync_with_stdio(False) # cout高速化
```

これでac-libraryや高速`input`、自作のC++ライブラリはPythonから呼び出せます(ただし`def`、`cpdef`を使用する必要があります)。

### 細かい指定
`O3`や`mavx512f`といった最適化指定は次のようにすることでできます。

```pyx
# distutils: extra_compile_args = -O3
```

また、`numpy`をcimportして使ったりac-libraryを使ったりしたい場合は`include_dirs`を追加します。

```pyx
# distutils: include_dirs = /home/contestant/.local/lib/python3.8/site-packages/numpy/core/include
```

これらは、`setup.py`を実行してコンパイルすることでも指定できます。

```py
import sys
if sys.argv[-1] == 'ONLINE_JUDGE':
  mycode = r'''
# distutils: language = c++
# cython: language_level = 3
def main():
  pass # ここに書く
if __name__ == 'mycode':
  main()
'''
  with open('mycode.pyx', 'w') as f:
    f.write(mycode)
  setup = r'''
from distutils.core import setup, Extension
from Cython.Build import cythonize
from numpy import get_include
ext = Extension(
  'mycode',
  sources = ['mycode.pyx'],
  include_dirs = ['.', get_include()], # これでnumpyが使える
  extra_compile_args = ['-O3', '-unroll-loops', '-mavx512f'] # これで最適化指定ができる
)
setup(name = 'mycode', ext_modules = cythonize([ext]))
'''
  with open('setup.py', 'w') as f:
    f.write(setup)
  import subprocess
  with open('error.txt', 'w') as ef:
    code = subprocess.run('python3.8 setup.py build_ext --inplace'.split(), stderr = ef).returncode
  if code:
    with open('error.txt', 'r') as ef:
      with open('mycode.py', 'w') as f:
        f.write(f'''
import sys
sys.stderr.write("""{ef.read()}""")
exit({code})
''')
  exit()
import mycode
```

以下は $O(N ^ 2)$ コードを通す定数倍高速化例です。

https://atcoder.jp/contests/abc091/submissions/25632759

Cythonのコンパイラは意外と優秀です。

### inline

```pyx
cdef inline void func():
  ...
```

でインライン化を指定します。

### キャスト

型を変換したい場合は以下のようにします。

```pyx
cdef int a = 100000000, b = 100000000
cdef long long c = <long long>a * b
```
### cdivision, nonecheck, wraparound, boundscheck

- `cdivision`: あまりがマイナスになることのあるC++の除算仕様の有無
- `nonecheck`: 関数の引数が`None`かのチェックの有無(関数の引数に`not None`をつけることでも同じ効果となる)
- `wraparound`: `list`型のマイナスindexの有無
- `boundscheck`: `list`型の範囲チェックの有無

`wraparound`, `boundscheck`は`vector`に関係ないので使わないです。

```pyx
# cython: cdivision = True, nonecheck = False
```

### for文の構文

```pyx
for 0 <= i < 1000 by 2: # 0, 2, 4, ..., 998
  ...
```

のようにすると、`i`を型指定していなくても最適化され高速になります。`by`は省略可です。

### その他の機能

enum, typeid, cython.parallel, public, readonly, except +などたくさんあります。

Cythonは奥が深いです。Cython公式ドキュメント https://cython.readthedocs.io/en/latest/ に全部書いてあります。

# 最後に

Cythonを学ぶことはC++の入門になります。

競プロにおいてはC++の方が使いやすいことも多くなってくるので、Cythonを使うことが多くなったらさっさとPythonからC++に乗り換えた方がいいと思います。(僕はそうしました)

最後まで読んでいただきありがとうございました。