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

はじめに

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

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

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

Cythonとは

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

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

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

Cythonの威力

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

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

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で提出します：

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の力を解放する魔術詠唱 を参考にしています。

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

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

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

静的型付け

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

基本構文

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

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

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

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

静的型付けできるもの

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

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

ctypedef

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

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

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

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

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

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$ で割った余りを求めるプログラムで試してみます。

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

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

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

# 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内にあり、以下が使用可能です。

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

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

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するといいです。

from libcpp.map cimport map as cmap

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

cdef vector[int] a

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

cdef string s = b'Hello world!'

高速化例

Educational DP Contest O - Matching を例にしてみます。

PyPyでは1300msほどです。

Cythonでは700msでした。

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

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

型推論

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

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

関数

基本構文

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

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

fused type

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

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

chmax関数を例に取ります。

# 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

基本構文

cdef class Cls:
  ...

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

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

# 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の構文上できません。そのため、ポインタを使うしか方法がないです。

# 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です。

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の中心的な機能の一つです。

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

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

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

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

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

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

# 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

# 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といった最適化指定は次のようにすることでできます。

# distutils: extra_compile_args = -O3

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

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

これらは、setup.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)$ コードを通す定数倍高速化例です。

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

inline

cdef inline void func():
  ...

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

キャスト

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

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に関係ないので使わないです。

# cython: cdivision = True, nonecheck = False

for文の構文

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++に乗り換えた方がいいと思います。(僕はそうしました)

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