例えば$10000\times 10000$の2次元配列全体に、何らかの処理をするとします。
2重forループが思い浮かぶところですが、pythonのforループが激遅いというのは有名です。
（正確には「配列インデックスアクセスが遅い」「基本的な演算に毎回型チェックが入るから遅い」などの複合的な要素があるのでやや語弊がある表現ではあります。）
pythonで単純な処理をループする場合の高速化の手段のひとつに、numpyがあります。
一方で、numpyで簡潔に書ける処理はある程度限定されており、全ての計算をnumpyで書くのは難しい場合もあります。そのときに使える考え方のひとつに、forループ1重分をnumpyにするというものがあります。具体例を見ていきましょう。
二項係数の生成
$0\leqq n,k\leqq 10000$に対して、二項係数${}_n\mathrm{C}_k$の値を生成しましょう。
値が非常に大きくなる（$\sum_{k=0}^n{}_n\mathrm{C}_k = 2^n$より、${}_{10000}\mathrm{C}_k$の最大値は$\frac{2^{10000}}{10001}$以上$2^{10000}$以下で、およそ$3000$桁の数となる）ので、ここでは下$10$桁を計算することにします。
二項係数の計算方法もいくつかの方法がありますが、今回は漸化式
\[
{}_n\mathrm{C}_k = {}_{n-1}\mathrm{C}_{k-1} + {}_{n-1}\mathrm{C}_k
\] を用いることとします。まずは2重ループによる実装です。
31秒かかりました。これをnumpyで高速化します。漸化式に注目すると、
comb[1:, 1:] = comb[:-1, 1:] + comb[:-1, :-1]という式が思い浮かびますが、これでは上手くいきませんね。comb[$n$, $k$]を計算する段階では、「計算結果の」comb[$n-1$, $k$]などの値を用いる必要があり、これらの計算を同時に行うことはできないからです。また、modもある程度の頻度で取ってあげないといけません。
そこで、それぞれの行ごとにはループ処理をしたまま、各行をnumpyで計算してあげます。
1秒で計算が終わっており、約30倍の高速化に成功しました。
前者のコードでは、「ループ」や、ループ変数による配列参照が$O(10^8)$回行われています。
一方で、後者のコードでは、それらは$O(10^4)$回しか行われていません。オーバーヘッドとなっていた部分の処理回数のオーダーを削減することで、総合的な処理時間を大幅に削減することに成功しています。
＜まとめ＞
★ 2次元ループ を 1次元ループに変更することで、ループに起因するボトルネックのほとんどを解消できる。
★ したがって大きな軸1方向についてnumpyが使えれば、巨大な2次元配列は十分高速に処理できる。
問題例
最近私がAtCoderで経験した例になります。ネタバレを含むのでご注意ください。
いずれも 2019/06/19 現在で、Python, PyPyでの正答提出者の中で最も実行時間の短いコードを実現しています。
AtCoder ABC 130-E Common Subsequence
次のような漸化式によって、dpおよびdp_cumを更新していく問題です。
\[
dp[n,m] = \begin{cases}dp_{cum}[n-1, m-1] & (S[n] = T[m])\\0 & (\text{otherwise})\end{cases},\qquad dp_{cum}[n,m] = \sum_{i\leq n, j\leq m}dp[n,m].
\] 全ての配列の要素をこのルールで一気に計算するのは難しそうですが、$n$ を固定して $m$ に関するベクトルだと思います。すると、$dp[n]$ は $dp[n-1]$ のシフト（スライス）です。ただし $T$ が特定の値となる場所を $0$ に書き直さなければいけませんが、これも簡単です（np.whereを使う、あるいは $(T == S[n])$ というbool値のarrayを掛ける。さらに $dp_{cum}[n]$ も、 $dp_{cum}[n-1]$ に $dp[n]$ の累積和を加えたものとして、numpyで容易に計算できます。
$n$ についてのループが残りますが、1次元分のループをnumpyで高速化してしまえば、そこはボトルネックには本質的にはひびきません。
実際の実装では、上述の $dp_{cum}$ に相当するものを $dp$ として処理を進めています。
AtCoder ABC 018-C 菱型カウント
AtCoder AGC 033-A Darker and Darker
この2問は、ほとんど同じ問題です。各マスが、最寄りの黒マスからどれだけの距離にあるかを調べる問題です。黒マスからの距離を、
・左にしか移動できないとして考えた数値
・左・右にしか移動できないとして考えた数値
・左・右・上に…
と更新していきます。各段階での更新では、
\[
dp[n,m] = \min(dp[n-1,m]+1, dp[n,m])
\]
という種類の漸化式を使うことになります。これを$dp[n]$というベクトルを$dp[n-1]$というベクトルから計算していると見なせば、1次元分のループをnumpy化することは容易です。
