はじめに
Python で競技プログラミングをしている人々の間では、「PyPy の文字列結合( += )は CPython のそれに比べてオーダーレベルで遅い」という罠はもはや常識になっています。
しかし、これは「PyPy の文字列結合が遅い」というわけではないのです。この間違った認識が広まっているせいで、PyPy が文字列処理に弱いようなイメージを持っている人も少なくないようです。1
この記事では、PyPy で文字列結合が O(N) かかる理由と、CPython だと O(1) な理由を説明するとともに、「CPython なら O(1) だから安心」という認識だとまずいぞという話をします。
ちなみに、記事内での計算量の表記に用いる N という文字は、ループ回数だったり文字列の長さだったりします。文脈から読み取ってください。
PyPy の += は遅いらしい
こんなコードを考えてみます。
res = ""
for i in range(10 ** 6):
res += "a"
aを10^6回繰り返した文字列を作る、なんの変哲もないコードです。( "a" * 10 ** 6 と書けるというのは置いておいて)
これを CPython で実行してみると、一瞬で終わります。10^6 回程度のループなら、いくら Python といえど余裕です。
しかし、PyPy だと一向に実行が終わりません。数十秒待ってやっと終了しました。
理由は、CPython はこのコードを O(N) で実行するが、PyPy は O(N^2) で実行するからです。
より具体的には、res += "a" の部分を O(1) で実行するか O(len(res)) で実行するかの差です。
では、PyPy はなぜこれを O(N) かけて実行するのでしょうか？
mutable と immutable
例えば、これが文字列ではなくリストへの連結だったら、当然 O(1) で連結してくれます。
res = []
for i in range(10 ** 6):
res += ["a"]
リストと文字列の違いは、mutable か immutableか(変更可能か不可能か)にあります。
リストは mutable ですから、+= されたら単にそのリストそのものの末尾に要素を足してしまえばよいです。リストは普通の配列として実装されていますから、あるリストの末尾に(定数個の)要素を追加するのは定数時間でできます。
一方で、文字列は immutable なので、その文字列自体を変更してしまうことができません。なので新しい文字列オブジェクトを生成し、その中身を結合結果とするのですが、この操作には当然文字列の長さに比例した時間がかかります。
# リストは mutable なので、オブジェクトそのものの中身を変更する実装になっていて、
# 単に代入して += するともとの変数の中身も変更される
a = [1, 2, 3]
b = a
b += [4]
print(a, b)
# => [1, 2, 3, 4] [1, 2, 3, 4]
# 文字列は immutable なので、
# += は新しい別のオブジェクトを生成して代入するという挙動になる
a = "abc"
b = a
b += "d"
print(a, b)
# => abc abcd
ということで、Python の仕様上、文字列の += が O(N) になるのはしょうがない、という結論でした。
…ん、じゃあなんで CPython は O(1) でできるんですか？
CPython の += が O(1) な理由
Pythonの仕様上、文字列の += は O(N) になってしまいます。しかし、最初に挙げたコードなどにおいては、CPython はこれを O(1) で実行します。
その理由はズバリ、「CPythonが訳のわからない最適化をしているから」 です。
参照カウント
まず、CPython のメモリ管理について軽く説明します。
メモリ管理というのは、つまり「要らなくなったオブジェクトを検知して消滅させる」ということをするための仕組みです。
CPython では、全てのオブジェクトについて、そのオブジェクトへの参照の数を管理していて、この数のことを参照カウントと呼びます。
オブジェクトへの参照というのは、例えば「変数にそのオブジェクトが入っている」「リストの要素にそのオブジェクトが入っている」「あるオブジェクトの属性にそのオブジェクトが設定されている」みたいなもので、変数にオブジェクトを代入したり、リストにオブジェクトを追加したりするたびに参照カウントが増えます。
逆に、変数に別のオブジェクトを代入すると、変数にもともと入っていたオブジェクトの参照カウントは1減りますし、リストから要素を削除すると、その要素の参照カウントは1減ります。
オブジェクトの参照カウントは0だと、「そのオブジェクトはもう必要ない」ということがわかり、オブジェクトが食っていたメモリを開放することができます。
そのオブジェクトへの参照がもう1つもないので、どう頑張ってもオブジェクトを使うことができませんから、当然ですね。
訳のわからない最適化
さて、天下り的ですが、ある文字列オブジェクトが入っている変数に += をするとき、そのオブジェクトの参照カウントが1ならば、そのオブジェクトそのものの中身を変更してしまっても問題がないことがわかります。
そのオブジェクトを参照する方法がその変数を介してしかないということですから、「immutableなはずのオブジェクトが知らん間に変更されてるんだが？」ということが起こらないんですね。
…つまり、CPython は、「+= の対象が文字列型のオブジェクトで、参照カウントが1であるときに限って、immutableな文字列オブジェクトの中身を変更してしまう」という最適化をすることで O(1) の文字列結合を実現しています。2
なので、参照カウントを適当に増やしてやると途端に O(N) かかるようになります。
res = ""
for i in range(10**6):
tmp = res # 適当な変数に res の中身のオブジェクトを束縛させることで、参照カウントを1増やす
res += "a" # この += には O(N) かかる！
この挙動、CPython を作っている人たちはどうとも思わないんでしょうか？
CPython3.11 だと O(N^2) かかる場合があるぞ！！！
ということで、CPython が訳のわからない最適化をしてくれるおかげで、文字列オブジェクトへの参照が変に増えないように実装し、間違えてPyPyで提出してしまわないように気をつけさえすれば、 += を高速にできて嬉しいですね。
…しかし、CPython3.11 以降ではそうとも行きません。CPython3.11 で冒頭のコードを実行してみると、全然実行が終わりません。
どうしてなんでしょうか。
CPython3.11 ではインタプリタに全体的な最適化が施され、その一環で二項演算をするためのバイトコード命令とその実装が変更されました。それにより、グローバルでのみ、この最適化が適用されなくなりました。ローカル、つまり関数内での文字列結合は、これまで通り最適化されます。
AtCoder で使用できる CPython のバージョンは 3.8 ですから、今のところ気にしなくても大丈夫です。
しかし、2023年の言語アップデートにおいて、CPython3.11 へのバージョンアップが予定されています。この言語アップデート以降は、CPython でもグローバルにおいては文字列結合に O(N) かかるようになります。
実際、Language Test 202301 のコードテストから冒頭のコードを実行してみるとめちゃくちゃ時間がかかります。
対処法としては、処理をmain関数に入れるとかになりますが、そもそも O(N^2) かかる書き方をやめて str.join() を使うようにするのが一番良いでしょう。
コンテスト中に罠を踏んでペナを付ける前に、推奨されない書き方を卒業しましょう。
おわりに
CPython3.12 以降でのさらなる最適化にあたって、グローバルでも最適化がかかる仕様に戻る可能性もあります。
もしそうなった場合、この issue にその旨のコメントがなされるでしょうから、最新情報はこちらを参照してください: https://github.com/python/cpython/issues/99862