🦁

AtCoderでnumpyを使いたい

2022/03/29に公開

自分用のメモに書いたつもりだったが、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=np.int)
b=np.ones((n,),dtype=np.int)
c=np.fill((n,),c0,dtype=np.int)
d=np.empty((n,),dtype=np.int)

要素数nの1次元配列の初期化例。zerosは初期値に0が代入される。onesは1,fullはc0が代入される。emptyは何が代入されるか分からないらしいがメモリ確保が最も高速。極限まで速さを追求するのでなければ基本zerosを使えば良いと思う。dtypeはデフォルトでは浮動小数点型。AtCoderでは多くの場合整数を扱うことが多いのでdtypeを指定する。AtCoder中でのnp.intは64bitの扱いになるので、普通に使う分にはオーバーフローしない。ただしpythonで使う整数と違い、2^64を超えるとnp.intではオーバーフローするはず。結構大きいサイズの配列でも、np.int64とnp.int8で計算速度はほとんど変わらないので、メモリ節約の目的がないのであればnp.intで良いと思う。実際の計測データだと巨大なデータを保持するのに必要なメモリが2GBか16GBでPCの要求スペックが変わるので、配列中の要素の最大値が256や2^16=65536未満であることが分かっているなら、np.int8やnp.int16を使う理由は十分ある。

座標圧縮もできる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オプションでヒストグラムが直接得られるので、出番は奪われがち。

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でも書いてくれる貴重なお方。

Discussion