🐌

AtCoder言語アップデート(2023年8月)後でもpypyでnumpyが使えない

2023/08/22に公開

AtCoderの言語アップデートによりPythonもバージョンアップ。これまではpypyではnumpyのimport時点でエラーが出てpypy+numpyの組み合わせは使えなかった。アップデートでpypyでもnumpyをサポートしてくれたので、これでnumpy使った解法の提出が増える、と思ったが・・・

アップデート後もpypyとnumpyの相性は悪く実用的でない

pypyでもnumpy, scipyのimportと関数が使えるようになったというだけで、numpyの速さが活きるかというとそんなことはなかった。
10^8の配列でコードテストしてみるとpypy+numpyの組み合わせはメリット無し。

python (cpython) + numpy

import numpy as np
n=10**8
a=np.arange(n)
a%=10
print(0)

10^8の配列の生成と余りの計算でだいたい500ms。

pypy + numpy

上記と同じコードをpypyで提出すると900~1000ms。2倍近く遅い。どうもnumpyのimportに時間がかかる模様。

import numpy as np
print(0)

このコードだけで500ms近くかかる。10^8の配列の演算で残り約500msかかる計算なので、numpy関数部分だけ見るとこのコードではpypyでもあまり実効速度変わらない。importが重い。

pypy + for loop

同じ演算をnumpy使わずに素直に書いてみる。

n=10**8
a=[0]*n
for i in range(n):
    a[i]=i
    a[i]%=10
print(0)

このコードだと500msで、python+numpyより気持ち早い。python+numpyはnumpyのimportに80ms程度かかるのでimportにかかる時間を抜いたnumpyとpypyは簡単な演算ではだいたい同程度の速度。

python (cpython) + for loop

上記のforループをpythonで提出すると10倍以上遅くなり9500ms前後かかる。pypyもnumpyも高速化の恩恵は大。pypy+numpyはpython+numpyより相対的に遅いが、numpy使わないpythonよりは圧倒的に早いので、まったくダメということはない。ただ、pypyの場合は無理にnumpyの関数使うよりはfor文で実装したほうが早そう、というのが結論。

pypy+numpyは何故遅い?

https://saturncloud.io/blog/why-is-pypy-slower-for-adding-numpy-arrays-a-deep-dive/
https://scipy.org/faq/
このあたりを読むと、numpyがC言語で最適化されている一方、pypyがpythonに特化して高速化を図っているので、Cで書かれたnumpy, scipyはpypyで最適化しにくいようだ。
もう少し複雑なnumpyでの処理をすると、pypy+numpyはpython+numpyより遅くなる傾向があった。importに時間がかかるだけでなく、numpyでの計算速度にも差が生じているようで、この速度差をnumpy以外の部分(forループ等)でpypyが逆転するのは厳しいように思う。そもそもforループ使わないようにnumpyで書くので。

numpy使った解法解説が欲しい・・・

欲しいが、おそらくこれまでと同様ABCのC問題以降はforループ駆使したpypy勢が多数派で、numpyでの解答はかなり限られると思われる。期待していただけに残念。

numpy派への朗報

numpyのimportにかかる時間がこれまでの約100msから70~80msに微妙に高速化した。また、scipyのバージョンアップによりDisjointSet(UnionFind)を呼び出せるようになった。ドキュメント見たところnumpyチックなベクトル処理は行わず要素一つ一つ処理する関数のようで速度の恩恵がどの程度あるか不明だが、コードは短くすっきりするかもしれない。これまでよく下記のサイトにお世話になってコードコピペしていたが、アップデート後は一行で済みそう。
https://note.nkmk.me/python-union-find/

from scipy.cluster.hierarchy import DisjointSet

Discussion