[numpy] 2次元配列の高速化

例えば $10000 \times 10000$ の2次元配列全体に、何らかの処理をするとします。
2重forループが思い浮かぶところですが、pythonのforループが激遅いというのは有名です。
（正確には「配列インデックスアクセスが遅い」「基本的な演算に毎回型チェックが入るから遅い」などの複合的な要素があるのでやや語弊がある表現ではあります。）

pythonで単純な処理をループする場合の高速化の手段のひとつに、numpyがあります。
一方で、numpyで簡潔に書ける処理はある程度限定されており、全ての計算をnumpyで書くのは難しい場合もあります。そのときに使える考え方のひとつに、forループ1重分をnumpyにするというものがあります。具体例を見ていきましょう。

二項係数の生成
問題例

二項係数の生成

$0 ≦ n, k ≦ 10000$ に対して、二項係数 $_{n} C_{k}$ の値を生成しましょう。
値が非常に大きくなる（ $\sum_{k = 0}^{n}_{n} C_{k} = 2^{n}$ より、 $_{10000} C_{k}$ の最大値は $\frac{2^{10000}}{10001}$ 以上 $2^{10000}$ 以下で、およそ $3000$ 桁の数となる）ので、ここでは下 $10$ 桁を計算することにします。

二項係数の計算方法もいくつかの方法がありますが、今回は漸化式
$_{n} C_{k} =_{n - 1} C_{k - 1} +_{n - 1} 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を更新していく問題です。
$d p [n, m] = {\begin{cases} d p_{c u m} [n - 1, m - 1] & (S [n] = T [m]) \\ 0 & (otherwise) \end{cases}, d p_{c u m} [n, m] = \sum_{i ≦ n, j ≦ m} d p [n, m] .$ 全ての配列の要素をこのルールで一気に計算するのは難しそうですが、 $n$ を固定して $m$ に関するベクトルだと思います。すると、 $d p [n]$ は $d p [n - 1]$ のシフト（スライス）です。ただし $T$ が特定の値となる場所を $0$ に書き直さなければいけませんが、これも簡単です（np.whereを使う、あるいは $(T == S [n])$ というbool値のarrayを掛ける。さらに $d p_{c u m} [n]$ も、 $d p_{c u m} [n - 1]$ に $d p [n]$ の累積和を加えたものとして、numpyで容易に計算できます。

$n$ についてのループが残りますが、1次元分のループをnumpyで高速化してしまえば、そこはボトルネックには本質的にはひびきません。

実際の実装では、上述の $d p_{c u m}$ に相当するものを $d p$ として処理を進めています。

AtCoder ABC 018-C 菱型カウント

問題文：■
自分の提出：■

AtCoder AGC 033-A Darker and Darker

問題文：■
自分の提出：■

この2問は、ほとんど同じ問題です。各マスが、最寄りの黒マスからどれだけの距離にあるかを調べる問題です。黒マスからの距離を、
・左にしか移動できないとして考えた数値
・左・右にしか移動できないとして考えた数値
・左・右・上に…

と更新していきます。各段階での更新では、
$d p [n, m] = min (d p [n - 1, m] + 1, d p [n, m])$
という種類の漸化式を使うことになります。これを $d p [n]$ というベクトルを $d p [n - 1]$ というベクトルから計算していると見なせば、1次元分のループをnumpy化することは容易です。