久しぶりにCythonを触る

AtCoderの新環境でCythonからC++の標準ライブラリが使えるようになり、より便利になった。Cythonはしばらく触っていなかったのだけど、いい機会なので再び触ってみることにした。

C++の標準ライブラリは、例えばvectorだとこんな感じで使える。

from libcpp.vector cimport vector

cdef vector[int] a
a.push_back(42)

気を付けなければいけないのは、以下のように書くとパフォーマンスが低下すること。

サンプルコードではグローバル直置きになっているが、ローカル変数については型推論が効く。まあ推論してくれるやろ～とこの書き方で書いたら推論してくれなくて困ったりした。

from libcpp.vector cimport vector

# 遅い
a = vector[int]()

# 型を指定したらセーフ
# cdef vector[int] a = vector[int]()

a.push_back(42)

そんな感じで、Cythonを使ってABC153-E "Crested Ibis vs Monster"を解くと以下のコードになる。solve()内の処理は全てC++にコンパイルされるように書いた。

# cython: language_level=3
# distutils: language=c++

from libcpp.vector cimport vector
from libcpp.pair cimport pair
from libc.stdio cimport scanf, printf


cdef solve():
    cdef int h, n, a, b, INF = 1000000000
    cdef int i

    scanf("%d %d", &h, &n)

    cdef vector[pair[int,int]] ab = vector[pair[int,int]](n)

    for i in range(n):
        scanf("%d %d", &a, &b)
        ab[i] = pair[int,int](a, b)

    cdef vector[int] dp = vector[int](h+1, INF)
    dp[0] = 0

    for i in range(h):
        if dp[i] == INF:
            continue
        for magic in ab:
            a, b = magic.first, magic.second
            dp[min(h, i+a)] = min(dp[min(h, i+a)], dp[i]+b)
    printf("%d\n", dp[h])

solve()

どうかしら。scanfやprintfを使ったりして、ここまで来たらもはやC++で書く方が楽な気もするが、ともかくCythonで書くとこんな感じになる。

全部こうやって書く必要は無くて、Python的なコードとC的なコードを混在させることができるのがCythonのウリなのでほどほどに...。

普段Pythonを書く時と同じ気持ちで、

for a, b in ab:
    ...

と書こうとしたが、これは遅い。生成されたコードを読んでみると、pairをPythonオブジェクトに変換し、それをunpackするという流れになっていた。pairを直接unpackすることは出来ないらしい。

そんな感じで、書くときに気を付けないと踏み抜いてしまう罠がそれなりにある。完全にCにコンパイルされるコードを書こうとするのなら、cython --annotate a.pyxで出力されるhtmlと睨めっこする必要があるだろう。

最初はPyPyで提出して、TLEしたらCythonに切り替えてちまちま最適化する...という使い方になるかしら。うーん、微妙な立ち位置。ギリギリを攻めるならCythonのような中間の言語でなく、最初からC++で書いた方が良いだろうしなあ...。