LoginSignup
15
17

More than 3 years have passed since last update.

@yH3PO4

Educational DP Contest(B,D,E,F,I,L)で学ぶnumpy高速化

Last updated at Posted at 2020-04-29

動的計画法(DP)の典型問題が集まっているEducational DP Contestですが、問題によってはTL制限がきつく、Python3で普通に提出すると想定解でもTLEしてしまうことがあります。
多くの場合PyPy3で提出する、また新ジャッジではおそらくnumbaを利用することでACできるようになると思われますが、この記事ではnumpyを用いた高速化テクニックを施すことによってACすることを目指していきます。

  • Python3.4.3での提出のうちTLEとなっているものが少ない問題については扱っていません。
  • 得られた漸化式をどう処理するかということに絞って扱います。なぜその漸化式でよいのかといった話はほかの記事を参照ください。drkenさんの記事がおすすめです(というか漸化式はほぼこの記事に準じています)。
  • 執筆にあたり、他の方のAC提出の数々を大いに参考にさせていただいております。

B - Frog 2

漸化式

$dp[i]=$「カエルが足場 $i$ へと移動するのに必要となる最小のコスト」として、

dp[i] = \min_{\max(0, i - K) \leq j < i} |H[i] - H[j]| + dp[j]

となるように$i$を$1$から$N$まで回します。ただし、$dp$テーブルの初期値は$dp[0]=0$、それ以外は十分巨大な値とします。
素のPythonで書くならこうですね。

for i in range(1, N):
    for j in range(max(0, i - K), i):
        dp[i] = min(dp[i], abs(H[i] - H[j]) + dp[j])

例えば$N=6,K=3$の場合、以下のように手続きは進行していきます。

\begin{align}
&…\\
&\text{chmin}(dp[3],|H[3] - H[2]| + dp[2]) \tag{dp[3]確定}\\
&\text{chmin}(dp[4],|H[4] - H[1]| + dp[1]) \\
&\text{chmin}(dp[4],|H[4] - H[2]| + dp[2]) \\
&\text{chmin}(dp[4],|H[4] - H[3]| + dp[3]) \tag{dp[4]確定}\\
&\text{chmin}(dp[5],|H[5] - H[2]| + dp[2]) \\
&\text{chmin}(dp[5],|H[5] - H[3]| + dp[3]) \\
&\text{chmin}(dp[5],|H[5] - H[4]| + dp[4]) \tag{dp[5]確定}\\
&\text{chmin}(dp[6],|H[6] - H[3]| + dp[3]) \\
&\text{chmin}(dp[6],|H[6] - H[4]| + dp[4]) \\
&\text{chmin}(dp[6],|H[6] - H[5]| + dp[5]) \tag{dp[6]確定}\\
&…
\end{align}

配列演算化

一般にnumpyは配列に対する一括処理を得意とします。本問では$dp[i]$を逐次計算するのではなく、連続部分列 $dp[i:i+K]$に対する処理を、$i$を$1$から$N$まで動かしつつ実行していく処理に変更します。

例えば先の例において、$dp[3]$の値が確定してから操作$\text{chmin}(dp[6],|H[6] - H[3]| + dp[3])$がなされるまでには、ある程度間があいています。これを、$dp[3]$が求まってすぐに操作$\text{chmin}(dp[6],|H[6] - H[3]| + dp[3])$を行うようなアルゴリズムに変更しても答えは変わりません。
すなわち、

\begin{align}
&…\\
&\text{chmin}(dp[4],|H[4] - H[1]| + dp[1]) \\
&…\\
&\text{chmin}(dp[3],|H[3] - H[2]| + dp[2]) \tag{dp[3]確定}\\
&\text{chmin}(dp[4],|H[4] - H[2]| + dp[2]) \\
&\text{chmin}(dp[5],|H[5] - H[2]| + dp[2]) \\
&\text{chmin}(dp[4],|H[4] - H[3]| + dp[3]) \tag{dp[4]確定}\\
&\text{chmin}(dp[5],|H[5] - H[3]| + dp[3]) \\
&\text{chmin}(dp[6],|H[6] - H[3]| + dp[3]) \\
&\text{chmin}(dp[5],|H[5] - H[4]| + dp[4]) \tag{dp[5]確定}\\
&\text{chmin}(dp[6],|H[6] - H[4]| + dp[4]) \\
&… \\
&\text{chmin}(dp[6],|H[6] - H[5]| + dp[5]) \tag{dp[6]確定}\\
&…
\end{align}

したがって、この新たな手続きを素のPythonで書くと、

for i in range(1, N):
    for j in range(i, i + K):
        dp[j] = min(dp[j], abs(H[j] - H[i]) + dp[i])

となります。このとき、$j$のループの中身は互いに独立しているので並列に処理することができて、

dp = np.full(N + K, 10 ** 10, dtype=np.int64)
dp[0] = 0
for i in range(1, N):
    dp[i:i + K] = np.minimum(dp[i:i + K],
                             np.abs(H[i:i + K] - H[i - 1]) + dp[i - 1])
print(dp[N - 1])  # 0-indexedに変換しておく

と書けます。np.minimumとnp.absはいずれもndarrayの各要素に対して処理を行いndarrayを返す関数です。
このように、numpy高速化においては並列処理を括り出すことが最も重要なテクニックとなります。

実装の際は$dp$テーブルの長さを$N+K$まで拡張しておかないとIndexErrorになることに留意してください。後ろの方は捨てることになります。

D - Knapsack 1

漸化式

$dp[i][w_{sum}] = $ 「$i-1$ 番目までの品物から重さが $w_{sum}$ を超えないように選んだときの、価値の総和の最大値」として、

dp[i + 1][w_{sum}] = 
\begin{cases}
dp[i][w_{sum}] &(w_{sum} - w[i] < 0)\\
\max(dp[i][w_{sum}], dp[i][w_{sum}- w[i]] + v[i]) &(otherwise)
\end{cases}

となるように、$i,w_{sum}$を$(0,0)$から$(N,W)$まで回して更新していきます。ただし、$dp$テーブルの初期値はすべて$0$です。
素のPythonで書くと：

for i in range(N):
    for x in range(W + 1):
        tmp = dp[i][x - w[i]] + v[i] if x - w[i] >= 0 else 0
        dp[i + 1][x] = max(dp[i][x], tmp)

配列演算化

この問題の場合、$i$を固定した時$w_{sum}$ごとの処理は既に互いに独立しており、そのまま配列演算に直すことができます。
唯一厄介なのが条件分岐ですが、これは以下のような配列を用意してあげることで処理できます。

tmp[i + 1][w_{sum}]=
\begin{cases}
dp[i][w_{sum}- w[i]] + v[i]　&(w_{sum} > w[i]) \\
0 &(otherwise)
\end{cases}

このような配列は、np.zerosで初期化してから$w_{sum} > w[i]$部分をスライスして処理すれば作ることができます。配列$tmp[i]$は一度ずつしか使わないので、$tmp[w_{sum}]$という形の一次元配列を繰りかえし生成するようにしてメモリを節約しましょう。

dp = np.zeros((N + 1, W + 1), dtype=np.int64)
for i in range(N):
    tmp = np.zeros(W + 1, dtype=int)
    tmp[w[i]:] = dp[i][:-w[i]] + v[i]
    dp[i + 1] = np.maximum(dp[i], tmp)
print(dp[N][W])

ちなみに$dp$テーブルを一次元にして破壊的に更新していくこともできます。

dp = np.zeros(W + 1, dtype=np.int64)
for i in range(N):
    tmp = np.zeros(W + 1, dtype=int)
    tmp[w[i]:] = d[:-w[i]] + v[i]
    dp = np.maximum(dp, tmp)
print(dp[W])

こっちの方が若干早いようですね。

E - Knapsack 2

漸化式

$dp[i][v_{sum}] =$「 $i-1$ 番目までの品物から価値が $v_{sum}$ 以上になるように選んだときの、重さの総和の最小値」として、

dp[i + 1][v_{sum}] = 
\begin{cases}
dp[i][v_{sum}] &(v_{sum} - v[i] < 0)\\
\min(dp[i][v_{sum}], dp[i][v_{sum}- v[i]] + w[i]) &(otherwise)
\end{cases}

となるように、$i,v_{sum}$を$(0,0)$から$(N,10^5)$まで回して更新していきます。ただし、$dp$テーブルの初期値は$dp[0][0]=0$、それ以外は十分巨大な値とします。

実際の答えは、「$dp[N][v_{sum}]$ の値が$W$ 以下となるような$v_{sum}$の値の最大値」となります。

配列演算化

これも先と同様にやります。

Mv = 10 ** 5
Mw = 10 ** 11
# いずれも十分巨大な値

dp= np.full((N + 1, Mv + 1), Mw, dtype=np.int64)
dp[0][0] = 0
for i in range(N):
    tmp = np.full(Mv + 1, Mw, dtype=int)
    tmp[v[i]:] = dp[i][:-v[i]] + w[i]
    dp[i + 1] = np.minimum(dp[i], tmp)

さっきのように$dp$テーブルを一次元にする書き方でもよいでしょう。
最後の答えは以下のように出すことができます。

print(np.max(np.where(dp[N] <= W, dp[N], 0).nonzero()))

まず、np.where(dp[N] <= W, dp[N], 0)で$dp[N][v_{sum}]$ の値が$W$ より大きくなるものを全て$0$に置き換えます。次にnonzero()によって、$dp[N]$のうち値が$0$でないところのインデックスが得られるので、この最大値をnp.maxでとれば答となります。

F - LCS

漸化式

$dp[i + 1][j + 1] =$「$s$の$i$文字目までと$t$ の$j$文字目まででの LCS の長さ」として、

dp[i + 1][j + 1] =
\begin{cases}
\max(dp[i][j + 1], dp[i + 1][j], dp[i][j] + 1)
& (s[i]=t[j]) \\
\max(dp[i][j + 1], dp[i + 1][j])
& (otherwise)
\end{cases}

となるように、$i,j$を$(0,0)$から$(\text{len}(s),\text{len}(t))$まで回して更新していきます。ただし、$dp$テーブルの初期値はすべて$0$です。
具体的にLCSを構成するパートはDPの部分に比べて十分高速1なので本記事では飛ばします。

配列演算化

まず、ここから条件分岐を削りたいという気持ちになります。仮に $equal[i][j] =$「$s$の$i$文字目と$t$ の$j$文字目が一致していれば$1$、そうでなければ$0$」という配列があれば、

dp[i + 1][j + 1] =
\max(dp[i][j + 1], dp[i + 1][j], dp[i][j] + equal[i][j])

とできます。
実は、このような配列equalはnumpyで次のようにして作ることができます。

s = np.array(list(input()))
t = np.array(list(input()))
equal = (s[:, None] == t[None, :])

equalのところの右辺は、s.reshape(len(s), 1) == t.reshape(1, len(t))に等価です。すなわち、$\text{len}(s) \times 1$の配列と$1 \times \text{len}(t)$の配列を作って比較していることになります。
配列equalの要素はbool値ですが、Trueは$1$、Falseは$0$にあたるのでそのまま先の漸化式に入れて使うことができます。

次に、配列に対する演算の形で書き換えます。各配列$dp[i]$ごとに処理する上で、$dp[i + 1][j + 1]$の式に$dp[i + 1][j]$が出てくるのは厄介です。そこで、$dp[i + 1][j]$を外します。

\begin{align}
dp[i + 1][j + 1] =
\max(& \max(dp[i][j + 1], dp[i][j] + equal[i][j]), \\ &dp[i + 1][j])
\end{align}

従って、$i$を固定したとき$dp[i + 1][j + 1]$は次のような手順で求めることができると分かります。

  1. 各$j$に対して$dp[i + 1][j + 1] = \max(dp[i][j + 1], dp[i][j] + equal[i][j])$とする。
  2. $dp[i + 1]$の累積maxをとってできた配列を新たに$dp[i + 1]$とする。

1.はnp.maximumにより、2.はnp.maximum.accumulateにより実現可能です。最終的な実装は以下の通りです。

s = np.array(list(input()))
t = np.array(list(input()))

equal = (s[:, None] == t[None, :])
ls, lt = len(s), len(t)
dp = np.zeros((ls + 1, lt + 1), dtype=int)

for i in range(ls):
    dp[i + 1][1:] = np.maximum(dp[i][1:], dp[i][:-1] + equal[i])
    dp[i + 1] = np.maximum.accumulate(dp[i + 1])

このように、累積〇〇の形にするのはnumpy高速化におけるもう一つの有用な手法です。

ところで、このコードでは$dp$テーブルでdtype=intにしていますが、これは要素が32bit整数に収まる場合にのみ利用可能です。怪しい時にはnp.int64にしておくのが無難です(ただしおそらく若干遅くなる)。

I - Coins

漸化式

$dp[i][j] = $ 「前から$i$枚のコインを投げたときに表になっているのが $j$ 枚ある確率」として、

dp[i + 1][j + 1] += dp[i][j] \cdot P[i] \\
dp[i + 1][j] += dp[i][j] \cdot (1 - P[i])

となるように、$i,j$を$(0,0)$から$(N-1,N-1)$まで回して更新していきます。ただし、$dp$テーブルの初期値は$dp[0][0] = 1$、それ以外はすべて$0$とします。

配列演算化

この問題に関しては特に式変形など経ることなくnumpy配列の演算に書き換えることができます。

dp = np.zeros((N + 1, N + 1), dtype=float)
for i in range(N):
    dp[i + 1][1:] += dp[i][:-1] * P[i]
    dp[i + 1][:-1] += dp[i][:-1] * (1 - P[i])
print(sum(dp[N][(N + 1) // 2:]))

$dp$テーブルでdtype=floatにするのを忘れないようにしてください。

J - Sushi (まだ)

2020/4/29時点でnumba不使用の提出が5本とかしかないです……

L - Deque

漸化式

$dp[l][r] = $「区間$[l,r)$が残った状態から最適に動いた場合の$X-Y$」として、

dp[l][r] = 
\begin{cases}
\min(dp[l + 1][r] - A[l], dp[l][r - 1] - A[r - 1]) &((N - (r - l)) \% 2 = 0) \\
\max(dp[l + 1][r] + A[l], dp[l][r - 1] + A[r - 1]) &(otherwise) \\
\end{cases}

となるように、$l,r$を$(0,1)$から$(0,N)$まで、$r-l$ が小さい方から大きい方へと回して更新していきます。ただし、$dp$テーブルの初期値はすべて$0$です。

配列演算化

for文の回し方がいつもと違っていて困るので、漸化式の方を書き換えます。$dp[m][l] = $「区間$[l,m+l)$が残った状態から最適に動いた場合の$X-Y$」として、

dp[m][l] = 
\begin{cases}
\min(dp[m - 1][l + 1] - A[l], dp[m - 1][l] - A[l + m - 1]) &((N - m) \% 2 = 0) \\
\max(dp[m - 1][l + 1] + A[l], dp[m - 1][l] + A[l + m - 1]) &(otherwise) \\
\end{cases}

となるように、$m,l$を$(1,0)$から$(N,0)$まで更新していく、という式にしてもアルゴリズムとしては同じですね($r$を$m-l$で書き換えただけです)。

$m$が固定されたとき、各$l$の処理は並列に行えるので、これは配列演算で書けます。

dp = np.zeros((N + 1, N + 1), dtype=np.int64)
for m in range(1, N + 1):
    if (N - m) % 2:
        dp[m][:N - m + 1] = np.minimum(dp[m - 1][1:N - m + 2] - A[:N - m + 1],
                                       dp[m - 1][:N - m + 1] - A[m - 1:])
    else:
        dp[m][:N - m + 1] = np.maximum(dp[m - 1][1:N - m + 2] + A[:N - m + 1],
                                       dp[m - 1][:N - m + 1] + A[m - 1:])
print(dp[N][0])

おわりに

つづく？

参考

  1. DPが$O(\text{len}(s) \cdot \text{len}(t))$、構成は$O(\text{len}(s)+\text{len}(t))$  ↩

15
17
2

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
15
17