AtCoderでnumpyを使いたい

自分用のメモに書いたつもりだったが、google検索で結構上位に出てしまい、他の人が書いた有用な記事が探しにくい。ZennのSEO対策が優秀すぎるのだろうか。なるべく有用な情報のみ残すようにしておこう。

numpyを使いたい理由とAtCoderでのnumpyの使われ方

自分の研究の中で実験計測データの解析にnumpyを使うことが増え、プログラミングの練習にと思いAtCoderに参加するようになった。問題は、Pythonの参加者はとても多いのに、numpy使う人がほとんどいないこと。Python遅いと思うならnumpy使え、とネット検索では出てくるが、速さが重要視されるAtCoderでなぜかnumpyは全く人気がない。他の人の解答を参考にしたいと思っても解答例がほとんどないのが実情。皆どうやって高速化しているかというと、C問題D問題以降はpypyでの解答が圧倒的に多い。たまにnumbaのJITを見るくらい。Python最速はAOT使った高速化のようで、pypyよりもぶっちぎりで速い。numpyも問題によってはかなりいい速度が出るのに、なぜ皆使わないんでしょうね。

numpyの注意点

• 配列サイズは頻繁に変化させてはいけない
  要素の追加・削除を頻繁に行う場合はpython組み込みのlist, setを使う。for文の中で毎回要素数の異なるnumpy配列を作っていると遅い。np.stack関数で複数のnumpy配列を結合する場合も、元の配列が巨大だとメモリの確保に時間がかかる？配列サイズが分かっているなら、最初から要素数の大きい配列を初期化して用意しておいた方が良い。
• 要素数の異なる多次元リストはnumpy配列にできない
  たまに要素数が一定でない問題に出くわす。どうしてもnumpyで処理したい場合はダミーの要素を追加してからnumpy配列にする等工夫がいる。たくさんのクエリを処理する問題で配列サイズが変わることが多く困る。後述のようにfor文を回さないことがnumpyのコツなので、クエリごとに答えると高速化が期待できない。クエリ先読みしてnumpy配列に格納しようとするとダミー要素が必要。
• 配列の一部要素のみを頻繁に指定する必要がある場合はlistのほうが良い
  大きな配列のうち一部の要素にしかアクセスしないような場合、numpyだと遅くなる。for文で顕在化。numpy使うときは配列全体の更新をするような計算を心がける。
• numpyとlistの変換を頻繁に行ってはいけない
  listのほうが早い処理があっても、numpyとlistの変換に時間がかかるため、for文の中で変えない方が良い。途中・最後に数回だけ変換するような処理を考える。numpyからlistへの変換はtolist()を使う。
• for文はなるべく使わない
  そもそもfor文を使わなくとも良いようにnumpy使っているので、numpy使う処理を10^5回行うようなfor文ならnumpy使わない方が良い。ただしどうしても使わないといけない場面は出てくる。2重、3重のforループが出てきたら、for文一つだけでもnumpyに置き換えられないか考えたほうが良い。numpy配列が10^4台でfor文が10^3回の繰り返しくらいなら、見た目10^7回の演算で不安だが意外と高速。
• 要素は数値に変換
  文字列もnumpy配列にできるが処理はあまり得意では無さそう。np.where等で文字列を数値に変換してから処理すること。True,Falseのbool型はnumpyでも高速。

よく使う関数

Sort系

a_sort=np.sort(a,kind="stable")
ind=np.argsort(a,kind="stable") #ind: index
a_sort=a[ind]
ind2=np.lexsort((c,b,a))

sortはsort済みの配列を、argsortはsortした添え字の順番を返す。順番を覚えておきたいときはargsortで添え字を出力しておく。デフォルトだと元の配列の順番が入れ替わりたまに意図しない結果になることがあるので、kind="stable"オプションで安定ソートを使った方が良い。このオプションで遅くなって困る、ということは基本的に無いと思うので、kind="stable"は毎回指定しておいた方が無難。
lexsortは複数の配列を優先順位決めてsortするときに使う。添え字を返すので、実際のsort済み配列を得るには、返された添え字配列を元のnumpy配列の添え字に指定する。後ろに指定した配列ほどsortの優先順位が高い。上の例だとa→b→cの順でsortする。lexsortはあまり早くない気がする。

二分探索 np.searchsorted

ind=np.searchsorted(a,b)

AtCoderの問題で出るような10^6程度の配列だと、python組み込みのbisectのほうが早いらしい。numpyの利点は、探索する数値に配列を指定できること。上の例ではbに数値ではなく配列を代入することで、それらが挿入されるaのindexを一度にまとめて処理できる。bに配列を使えるとnumpyで処理した感が出る。

配列初期化

a=np.zeros((n,),dtype="i8")
b=np.ones((n,),dtype=np.int64)
c=np.fill((n,),c0,dtype=np.int64)
d=np.empty((n,),dtype=np.int64)

要素数nの1次元配列の初期化例。zerosは初期値に0が代入される。onesは1,fullはc0が代入される。emptyは何が代入されるか分からないらしいがメモリ確保が最も高速。極限まで速さを追求するのでなければ基本zerosを使えば良いと思う。dtypeはデフォルトでは浮動小数点型。AtCoderでは多くの場合整数を扱うことが多いのでdtypeを指定する。以前のAtCoderのシステムではdtype=np.intとすると64bit整数として扱ってくれた。2023年の言語アップデート後はdtype=np.intはエラーとなるので、陽に型を指定する必要がある。文字列として"i8"とすると符号付き64bit整数となる。np.int64としても同じなので好みの書き方でok。ただし、Pythonの整数と違い2^64を超えるとオーバーフローする。大きい数を扱う場合は上限を超えないか注意が必要。dtypeについては結構大きいサイズの配列でも、np.int64とnp.int8で計算速度はほとんど変わらないので、メモリ節約の目的がないのであればnp.int64で良いと思う。実際の計測データだと巨大なデータを保持するのに必要なメモリが2GBか16GBでPCの要求スペックが変わるので、配列中の要素の最大値が256や2^16=65536未満であることが分かっているなら、np.int8やnp.int16を使う理由は十分ある。ちなみに整数は64bitまでしかdtypeの用意がない。小数は特に指定しないと64bitの浮動小数点型(dtype="f8"かdtype=np.float64)になるが、小数は128bit(dtype=np.float128)まで使える。64bit浮動小数点の有効数字は16桁前後で、64bit整数で扱える18桁の整数は表現できない。128bit浮動小数点だと有効数字が30桁程度あるようなので、整数問題で使うにはちょっと心配はあるが、16桁以上の精度が必要で整数問題として解くのではなく小数で押し通したい場合は128bit小数が必要。

座標圧縮もできるnp.unique

配列中の要素の出現回数を数えるのに使う関数。return_counts, return_index, return_inverse等いくつかオプションがあり毎回google検索している。

u,inverse=np.unique(a,return_inverse=True)

とすると、inverseは元の配列を座標圧縮した配列になっているはず。

ヒストグラム取得に使えるnp.bincount

配列の要素のヒストグラムを得るnp.bincountはnumpyの中でもかなり高速。知っておいて損はない。下限が0で1刻みのヒストグラムになるので、使えるのは配列中の要素最大値が10^6程度まで。10^9はメモリに乗らない。座標圧縮すると適用できるが、座標圧縮するときに使うnp.unique関数のreturn_countsオプションでヒストグラムが直接得られるので、出番は奪われがち。

Setの感覚で使えるnp.isin

第一の配列の要素が第二の配列の要素に含まれるかどうかを判定する関数。第二の配列サイズに比例した計算時間だと勝手に勘違いしていたが、実際に使ってみると恐ろしく高速。たぶんO(1)の計算時間で、Pythonのsetの代わりとして使える。Pythonはlistとsetを区別しないといけないが、np.isin関数はいちいちsetに変換しなくてもよい。また、Pythonだと第一の配列の要素の判定にfor文を回す必要があるが、isin関数は一発でベクトル処理してくれる。複数の要素が別の配列の要素に含まれるか判定するときはPythonのsetよりも圧倒的にisin関数が早い。この関数の有用性を知ってから、Pythonのsetを使うことが減った。

numpy解答の探し方

他の人の解答からnumpyを使った解法を見つけるのは結構大変。使う人が少ないのでかなり見つけづらい。「すべての提出」言語タブでpython3を選ぶとpypyとpythonの解答が出てくるが、「言語」列（タブではない）からpythonを選ぶとpypyを除外できる(pypyではnumpy使えないので除外してしまって問題ない)。ただし圧倒的にpypyが多いのでpythonが一覧に無いことも多い。pypyと比べてpythonのほうがメモリ使用量が小さい傾向にあるので、提出時刻順でなくメモリ順に並べ替えるとpython解答が出てきやすい。numpyの特徴は速さとメモリ使用量に出やすいので、pythonの他の解答と比べて実行時間が早い、同程度の実行時間の中でメモリ使用量が大きく違う、という解答を探すとnumpy解答に巡り合う確率が上がる。人で探す場合は現時点ではAkari__さんの解答がスマートかつ爆速のことが多いので参考にしましょう。C～E,F問題でpython/pypy含めて圧倒的な実行時間での解答例が多い。numpyというよりはAOTでの高速化のことが多いが、たまにnumpyでも書いてくれる貴重なお方。