先日出題されたこちらの問題について
頑張って高速化を行ったので紹介したいと思います。
上記問題は普通に解くと3次元のDPテーブルを作成して更新していくことになります。
しかしPythonあるあるですが、多次元のDPテーブルを作成して問題を解こうとすると、多くの場合TLEします。
PyPyを使うことである程度改善可能ですが、この問題のように3次元かつ3 * 3000 * 3000で計算量自体も多くなる問題ではPyPyですらTLEする可能性が高くなります。
また実行時間もそうですが、次元が増えると添え字でバグらせる確率が高くなります。（少なくとも私はそうです。 2次元でぎりぎり、3次元テーブルの遷移を頭で追うのは無理です）
回避策としてはDPテーブルを使いまわして次元を減らすのですが、2次元=>1次元や、3次元=>2次元のように、1次元だけ落とすのはよくあるのですが、今回の問題では3次元から1次元まで落とすことができました。
順を追って解説を書いていきます。（問題文に出てくるR, Cが好きになれないので、h, wで置き換えます。hが高さ、wが幅です）
注意&謝罪 私は数学や情報科学の出身ではなく、厳密な証明や正確な記法がさっぱりわかっていません。普段の問題も何となくの雰囲気で解いているので、ふんわりしたイメージの表現が多くなっています。何卒ご容赦ください。
また、Y軸を考えるときに下から上に上がっていく人と、上から下に降りていくひとがいると思いますが、私は多くの場合後者なので、そのような表記になっています。適宜脳内変換してください。
まずはh=1, k=1
横1列にマスが並んでいて、マスにアイテムが置かれていて、そのうち1個だけ取るのを考えます。
この場合ただの貪欲で解けて、左から配列を見ていって、それまでの最高値で更新するだけです。
numpyでやる場合はmaximum.accumulateを使うと、スッキリ書けます。
# 8マスに以下のようにアイテムが配置されているとする(数値はアイテムの価値)
[0, 3, 1, 2, 6, 1, 8, 0]
# DPを各マスでの最高値とする(ただし0マス目も存在するとしてDPは長さ9で準備)
[0, 0, 3, 3, 3, 6, 6, 8, 8]
＃コードで書くと
items = [0, 3, 1, 2, 6, 1, 8, 0]
DP = np.zeros(len(items)+1, dtype=np.int64)
DP[1:] += items
DP = np.maximum.accumulate(DP)kを増やす
実際の問題ではk=3なので、複数個取れるようにします。
上で作成した配列はアイテムを1個だけ取る場合でした。
しかし、アイテムを2個取れる場合でもそれほど変化はありません。
あるマスに到着する直前の最高値に(つまりDPだと一つ前のインデックス)、そのマスで取れるアイテムの価値を足すと、2個目のアイテムをその場所で取った時の最高値になります。
items = [0, 3, 1, 2, 6, 1, 8, 0]
DP = np.zeros(len(items)+1, dtype=np.int64)
DP[1:] += items
DP = np.maximum.accumulate(DP)
# 8マスに以下のようにアイテムが配置されているとする(数値はアイテムの価値)
[0, 3, 1, 2, 6, 1, 8, 0]
# DPを各マスでの最高値とする(ただし0マス目も存在するとしてDPは長さ9で準備)
[0, 0, 3, 3, 3, 6, 6, 8, 8]
# 直前の最高値にそのマスのアイテムを足す
DP[1:] = np.maximum(DP[:-1] + items, DP[1:])
# こうなる
[0, 0, 3, 4, 5, 9, 7, 14, 8]
# 自分のマスより左で高い値があるなら、もちろんそっちの方がいいので更新する
DP = np.maximum.accumulate(DP)
# こうなる
[0, 0, 3, 4, 5, 9, 9, 14, 14]同じ操作を繰り返せばk=3, 4と増やせます。
items = [0, 3, 1, 2, 6, 1, 8, 0]
DP = np.zeros(len(items)+1, dtype=np.int64)
k = 3
for _ in range(k):
DP[1:] = np.maximum(DP[:-1] + items, DP[1:])
DP = np.maximum.accumulate(DP)DP[x軸][アイテムの個数]だったのがDP[x軸]になるのでこれで1次元減りました！！
どうでしょう。個人的にはこれくらいならなんとかDPの遷移を頭で追えるレベルです。
hを増やす(縦に移動する)
例えば4行目から5行目の移動を考えるときに、1行目や3行目でどのような値だったかの情報は必要ないです。（4行目までの最高値さえあればよい）
ということはDP[y軸][x軸]を、DP[x軸]を使い回すことで表現できそうです。これはよくやる手法ですね。
1行目から2行目に縦に移動する場合の処理を考えます。
まず最初は1行目の最大値の状態から真下に降りてきて、そこでアイテムを取った時を考えます。（例えば、y=1, x=1で下に降りたあと、何も取らずに右にy=2, x=4まで進んでそこでアイテムを取るというムーブはy=1, x=4で下に降りて y=2, x=4のアイテムを取るムーブよりよくなることは絶対にないのでこれで大丈夫なはずです。）
つまりDP各マス(0マス相当のDP[0]を除く)に2行目のアイテムの値を足すだけです。
繰り返しになりますが自分より左で高い値があるならそちらを採用すべきなので、同じように左から貪欲して更新します。
# 1行目のアイテムが下記とする
items_1 = [0, 3, 1, 2, 6, 1, 8, 0]
# 1行目だけ考えて3個までの最高値を求める
DP = np.zeros(len(items_1)+1, dtype=np.int64)
k = 3
for _ in range(k):
DP[1:] = DP[:-1] + items_1
DP = np.maximum.accumulate(DP)
items_2 = [#]
# 1行目の最大値でそのまま降りてきてアイテムを取ったとする(v=0もアイテムとする)
DP[1:] += items_2
# 自分より左で高い値があるなら(略
DP = DP.maximum.accumulate(DP)今の状態は2行目でアイテムを1個だけ取った時の各マスの最大値です。あとは1行目の時と同様にアイテムを取る数を増やすだけです。
# 今の気持ちは2行目でアイテム1個だけ取った時の最高値
# アイテムを増やすだけなので1行目と同じ操作を2回やる
# ただし、1行目と違って左のマスからやってくる値よりも
# 上からやってきたすでにある値の方が大きい可能性もあるのでmaximumで更新する
for _ in range(k-1):
DP[1:] = np.maximum(DP[:-1] + items_1, DP[1:])
DP = np.maximum.accumulate(DP)あとは同じ操作を各行で実行するだけです。1行目だけforループ3回であとは2回だとなんだかダサいので、0行目があるという気持ちで、以下のように書けばスッキリする気がします。
item_table = np.array([
[items_1],
[items_2],
[items_3], # 以下略
])
w = len(item_table[0])
DP = np.zeros(w+1, dtype=np.int64))
for items in item_table:
DP[1:] += items # 上から降りてくる
DP = DP.maximum.accumulate(DP) # 自分より左(略
for _ in range(2): # あと2個取れるので2回更新
DP[1:] = np.maximum(DP[:-1] + items_1, DP[1:])
DP = np.maximum.accumulate(DP)
print(DP[-1])どうでしょうか？？個人的にはとてもすっきり書けて高速&省メモリ化が達成できたのではと思っています。（なので記事にしてみました）
もっとスッキリかけるぜ！という方ぜひコメントください！！
ちなみに実行時間は536msで同一人物の提出を除けば今の所Pythonで第3位です。1, 2位はnumbaも使っているので（というか1~20位までみた範囲全員numba）、上記の実装をベースにnumbaを使えばダントツ1位をとれるかもしれません！（私は使い方がよくわからないのでやっていません・・・だれかやってください・・・）
以下は私の提出です。
コメント