LoginSignup
3
5
@AkariLuminous(Akari)

[Python] ユーザー定義クラスを載せても速いセグメント木をつくる with ChatGPT

Posted at

0. はじめに

セグメント木(SegTree)は、区間の集計や更新を効率的に行うためのデータ構造で、競技プログラミングにおいてしばしば出題されます。
AtCoderが提供しているライブラリAtCoder Library(ACL)にも入っており C++ 以外の言語への移植も行われているため、AtCoder で開かれるコンテストでは参加者は手軽にセグメント木を使用することができます。

ACL の SegTree は汎用性の高い実装になっており、ユーザーが「要素の型」、「要素間の演算」、「単位元」を指定して使います。Pyhon版では型の指定は不要で以下のように使用します。

class S:
    def __init__(self, value1, value2):
        self.value1 = value1
        self.value2 = value2

def op(a, b):  # 演算
    return S(a.value1 + b.value1, a.value2 * b.value2)

e = S(0, 1)  # 単位元


n = 5
seg = SegTree(op, e, n)

しかし Python では実行速度の問題から要素にユーザー定義のクラスを使うことは敬遠されがちで、値を並べたタプルや複数の値を １ つの値にエンコードしたもので代用されることが多いです。

例として ABC357 F - Two Sequence Queriesを見てみます。
この問題は遅延評価機能付きのセグメント木である LazySegTree を使う問題で実行時間制限が $5 \rm{sec}$ と長めに設定されています。
提出時点で CPython/PyPy で AC している人のうち $1000 \rm{ms}$ を切っているのは CPython が 2 人と PyPy が 1 人、 $2000 \rm{ms}$ でも PyPy が 7 人であり、多くは $3500 \rm{ms}$ 以上かかっていることからもこの問題の実行時間制限が厳しめであることが伺えます。
また、実装をいくつかみてみると、そのほとんどが要素をタプルとして保持していることがわかります。

そこで、公式解説のように要素をクラスとして定義して公式に導入されている Python 版 ACL を使って提出したところ CPython、PyPy ともに TLE でした。

もちろん、ユーザー定義のクラスを載せても AC している提出はあるので不可能というわけではありませんが、かなり厳しいことが予想されます。
さらに、この問題は答えを $998244353$ で割ったあまりを出力することが求められているので、あまりを自動で取ってくれるクラス modint を使った場合にはより遅くなります。

タプルで困ることがあるかと言われればあまり思いつきませんが、隣の芝生は青く見えるのでクラスを載せても十分高速な SegTree/LazySegTree を作ってみました。

今回は ChatGPT とペアプログラミングをしたので、最後に所感を載せたいと思います。

1. どんなものを作りたいか

まず、どんなものを作りたいかをまとめておきます。
大まかには次のようになります。

用法がACLに近い

使うために特殊な操作や記法がたくさん必要だと体験としてよくないので、できるだけ簡単に使えることが望ましいです。

問題ごとに自分で書く部分は Python のみ

内部で C++ が動いていたとしても貼るだけならば問題ありません。自分で書く部分(問題固有の部分)は全て Python で記述できるようにしたいです。

できるだけ高速に

指標としては ABC357 F を $3000 \rm{ms}$ 以内で AC できるくらい高速なものを作りたいです。

2. どうやって作るか

高速化のために

素の CPython/PyPy では速度が物足りないことはわかっているので、他言語の力を借ります。
すぐに思いつく使えそうな方法として

  • NumPy + Numba
  • Python/C API
  • Cython

がありますが、Numba はクラスの最適化がいまいちで、Cython は直近の AtCoder の言語アップデートでなぜか CPython から削除されていたので使えません。
なので C++ で書いた SegTree を API を通して Python で使えるようにします。

どこまで C++ 化するか

高速化のためにはできるだけ Python を避けることが重要です。例えば Python のクラスをそのまま載せられるようにすると、内部の配列や添え字の操作は C++ で動きますが、結局 「Python のクラス同士の演算を Python の関数が行う」ということが多数回行われるのであまり高速化できないことが予想されます。
よって、クラスも関数も「内部と呼べる部分は全て C++ で動く」ことを目指します。

Python 以外を書かないために必要なこと

型や演算も C++ で動いて欲しいですが、先ほど述べたように、これらの問題固有の部分は Python で記述したいです。
これを実現するために

  • Python のクラス、関数に対応する C++ の構造体、関数
  • Python のクラスと C++ の構造体をつなぐ変換装置

を自動生成するようにします。

どうやって自動生成するか

Python は C++ と違い動的型付け言語なので変数の型の宣言をしなくてもいいし、途中で型が変わっても問題ありません。しかしこの差分を吸収するのは大変そうなので、Python 側に幾つかの制限を設けることにします。

  • クラスのコンストラクタの引数や関数の引数で変数の型を annotation により宣言する
  • 途中で型を変更しない

Python -> C++ の変換はastモジュールを使います。
引数の型を annotation から取得しそれを伝播させていくことで新たに定義された変数の型を推定します。

全体のイメージ

segtree_1.png

3. 実際に提出してみる

作ったものを提出してみました。具体的なコードは長いので提出をみてください。
3640行目からが問題固有の部分です。

最初に例として出した問題以外にもいくつか提出してみました。
結果は以下のようになりました。Range Affine Range Sum 以外は提出時点における CPython/PyPy の AC のなかで実行時間順にソートした場合に 1 ページ目に入ることができました。(Range Affine Range Sum は 2 ページ目でした。)
CPython は PyPy に比べて for 文などで遅くなりがちなことを考えるとかなり高速化できたと言えると思います。

コンテスト名 問題 使うもの 提出 実行時間
ABC357 F - Two Sequence Queries LazySegTree
ModInt		 提出 $974\rm{ms}$
ABC322 F - Vacation Query LazySegTree 提出 $867\rm{ms}$
ACL Begginer Contest D - Flat Subsequenc SegTree 提出 $184\rm{ms}$
ACL Practice Contest K - Range Affine Range Sum LazySegTree
ModInt		 提出 $1697\rm{ms}$

4. 注意

このような C拡張を埋め込む方法はコンテストサイトによってはルール違反となりアカウントBANなどの措置が取られることがあります。特に、Codeforces では禁じられているので注意してください。(ルール)
どのコンテストサイトであっても提出する前にコンテストルールを確認してください。
また、紹介したものはバグを含む可能性がありますので使用する場合は自己責任でお願いします。(バグがあったら報告いただけると嬉しいです。)

5. ChatGPT をプログラミングに使ってみた所感

かなり良かった

結論から言うと ChatGPT を使って良かったです。
これまでも ChatGPT を使ってプログラミングをしたことは何度もあったのですが、うまくいくことも全くうまくいかないこともあり、今回どうなるか半信半疑でした。
今回のケースに近い具体的な例としては

  • 「AtCoder 上で動くユーザースクリプトを書いて欲しい」-> 成功
  • 「Numba で高速化できる形で木の DFS を書いて欲しい」 -> 失敗(エラーが解消できず)

といった感じでした。
今回のケースでは Python/C API の使い方や Python -> C++ の変換といった部分は学習データが少なそうなので心配だったのですが、想像以上にうまくいって驚きました。

4o と 3 の違いはあまり感じず

私は無課金ユーザーなのでコードを出力させるとすぐに上限に達してしまい、3 に切り替わっていました。ですが、3 になったからといって急に精度が落ちたといったことは感じられませんでした。むしろ生成が速いことが多く、プログラミングという目的に限っては 3 の方がいい場合もあるのかもしれません。

今回の目的に合っていた

今回つくったものは

  • (今回は良かったけど)実はうまくいかない
  • 実はまずいことをやっている(参照カウントとか)
  • 保守が大変

といったことがあったとしても期待した出力が得られさえすればいいものでした。
このような目的には特に ChatGPT 等のツールを使ったプログラミングが適していると感じました。

完全に ChatGPT に任せた方がいい？

基本は

  1. 「こういうものを作って」と要望を伝える
  2. 出力されたコードを手元で実行して確認する
  3. エラーが出た場合はそれを伝えて修正してもらう

といった流れを繰り返していたのですが、

  • 式を見やすく整形したい
  • 変数名等を変更したい
  • ちょっとした修正をしたい

のようなことを手元で行って ChatGPT の保持する情報をアップデートするのを忘れるといったことが多発しました。
こういった自分でやった方が楽なものも ChatGPT でやった方がいいのでしょうか。それとも、変更するたびにコピーして ChatGPT に渡すのを意識づける方がいいのでしょうか。
上手い使い方が知りたいです。

6. おわりに

ユーザー定義クラスを載せても速いセグ木を作ってみました。予想よりも高速化できたのが嬉しい誤算でした。

もう少し拡張すれば別のモノイドを持つ複数のセグ木を作れそうですし、応用すればstd::setやstd::multisetも使えるようになりそうです。いつかやってみるかもしれません。

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

3
5
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
3
5