ABC438D - Tail of Snake
https://atcoder.jp/contests/abc438/tasks/abc438_d
多くの人に解かれている問題なのに僕には解けませんでした。どうやったら解法が思い浮かぶのか自分なりに考えました。
コンテスト中の動き
19分でC問題までをAC。A, B ともにちょっと考えさせられました。今回は少し難しかった気がします。Cはよく見かけるパターンでした。そして D 問題。8分ほど考えてわからなかったので E 問題を見ました。周期とかを考えて計算すれば答えの出る問題なのでいけるだろうと思って実装に取り組みました。ところが……40分ほど使った挙げ句諦めるという最悪の結果になりました。今思えばここで欲を出したのがまずかったです。
残り36分ほどで D 問題に戻り、再び考え始めましたが良いアイデアは出ません。結局貪欲法を考えて書きましたが答えは合いませんでした。コードの完成度にも自信がなく、どこが悪いのか考えている内にタイムアップです。およそ3ヶ月ぶりの茶パフォーマンスでレートはまたも激落ちです。
考察
要するに x, y の位置を動かして最適な位置を決める問題ですね。
- 全探索は最悪 $O(N^2)$ の計算をするはめになって言うまでもなく TLE。
- 3つの数列を使うということで、中段を固定して上下段を動かすやつ？ → 固定の段階で x, y を決めるのに最悪 $O(N^2)$。
- x, y の一方を固定して他方を動かす。x を固定すれば A が固定され、あとは B / Cのどこを区切るかの問題になる。その後は
- 二分探索 → 単調性がないので無理。尺取法も同様。
- 数式にしてみてなんとかして $ B+C $ が単調性を持つような式は立てられないか？ → 無理だった。
ここまで考えて無理だったので諦めて E 問題に行きました。そして復帰後、少ない時間でとりあえず考えた方針です。
- B / C の区切りを1つ動かすとその差分だけ値が上下する。差分を計算した後累積和を取ることで、どの位置が最も高い値になるのかがわかる。
- A / B の区切りを動かしながらその最も高いところを求める。区切った位置 x よりも右にある中で最も高い値を取るようにする。
結局これではダメでした。コンテストの後でゆっくり考えてわかりましたが、手順 1. でとった累積和は x が動いて B + C の取り得る範囲が変わったら無効になりますね。x より左の B, C の中にあった値が無視されるので、当然それより右で計算した累積和もそのままの値では使えないわけです。
時間が残り少なく焦っていたこともあり方針を疑う余裕はありませんでした。また、自分が書いたコードにバグが潜んでいる可能性をずっと考えていました。こんな解法は今までにやったことがなく、添字のズレ等でどこに間違いがあるか想像もつかなかったからです。後から思えば ACx10 / WAx28 という結果だったのでコーナーケース等ではなく冷静に方針そのものがおかしい可能性を考えるべきだったのかもしれません。ただまあ最後の手段とも思っていたのでそれは難しかったかもしれません。
解答
さて、正しい解法は公式解説を読めばわかりますが DP （動的計画法）です。可能性を考えもしなかったので驚きました。左から i 個目のブロックを j にしたとき（A, B, C をそれぞれ 0, 1, 2 とする）の「らしさ」の合計値の最大値を dp[i][j] に格納するようにします。
実装
N = int(input())
A = [0] + list(map(int, input().split()))
B = [0] + list(map(int, input().split()))
C = [0] + list(map(int, input().split()))
# dp[i][j] : i 番目に j = A/B/C を選んだときの「らしさ」の合計値の最大値
dp = [[-1 for _ in range(3)] for _ in range(N+1)]
dp[1][0] = A[1]
for i in range(1, N): # 配る DP
for j in range(3):
if dp[i][j] == -1:
continue
if j == 0: # "A" の次は "A" もしくは "B"
dp[i+1][0] = max(dp[i+1][0], dp[i][0] + A[i+1])
dp[i+1][1] = max(dp[i+1][1], dp[i][0] + B[i+1])
elif j == 1: # "B" の次は "B" もしくは "C"
dp[i+1][1] = max(dp[i+1][1], dp[i][1] + B[i+1])
dp[i+1][2] = max(dp[i+1][2], dp[i][1] + C[i+1])
elif j == 2: # "C" の次は必ず "C"
dp[i+1][2] = max(dp[i+1][2], dp[i][2] + C[i+1])
print(dp[N][2]) # 最後は必ず C を選んでいるため
まとめ
https://qiita.com/omakasessan/items/779ceafbd945db480a90
前回の C 問題は DP と思いきや貪欲法でした。一方、今回の D 問題は貪欲法と思いきや DP でした。翻弄されています。自分の中でもこれらの区別ができずにいますが、一つの方針としてはまず DP や他の解法を検討して最後に貪欲法を試すというのがあります。貪欲法は反例を挙げるのが難しいことが多いからです。
今回は DP そのものが思いつかなかったのでどうにもなりませんでした。最近は解法は簡単なのにそれを思いつくまでが大変で他の解法と惑いやすい問題が本当に増えた気がします。ですが多くの人は惑わされずに答えにたどり着けているので、これは僕が環境の変化に追いつけていないということなります。どういう勉強をすればこれが治せるのかわからずにもがいている状態です。あまり歳のせいにはしたくないものですが。
貪欲か DP か？
なぜ DP を思いつきもしなかったのかを考えていたのですが、考察のところで挙げた解法からもわかるように僕は「仕切りを動かす」ことしか考えていませんでした。よくよく考えてみるとこの ABC438D はこういう問題ですね。
- 各列 $ 1, ..., N $ において A / B / C を選ぶ
- 一つ前にどれを選んだかによって次に選べる英文字が制約を受ける
これなら僕でも DP とわかります。この問題、僕は仕切りを左右に動かすことしか考えていませんでした。しかし実際には要素を上下に動かす問題でした。図でいうと上ではなく下のイメージです。
これは僕の感覚的な話なので共感はされなくていいんですがこんな感じですね。困ったらこれを思い浮かべて別の解法が思い浮かぶようにしておきたいものです。
| DP | 貪欲その他 |
|---|---|
| 要素が主役 | 境界（区間）が主役 |
| 縦に動かす | 横に動かす |