はじめに
- ABC370のe問題は本番では解けなかったけど、解説の冒頭で「dpを使う」というの見て、なるほどと思って結構楽勝で解くことが出来た...TLEが混ざってたけどね！
- なので、解説をちゃんと読んで、TLEが出ないような解き方を理解したので、紹介したい。
どんな問題？
- サイズNの整数配列に区切りを入れる（部分配列を作る）方法はいくつあるか？ただし、全ての部分配列で合計がKにならないようにする。
- 例えば、サイズ3の配列[3 5 1]に区切りを入れる方法は4つある、すなわち「3/5/1」「3/5 1」「3 5/1」「3 5 1」の4つ。サイズNの場合は、$2^{N-1}$となる。
- このとき、「部分配列の合計がK=6になってはならない」という制約を適用すると、[3 5 1]の場合に可能な区切りを入れる方法は3つとなる。「3/5 1」の部分配列合計が6になってるから除く必要がある。
- 解答は「mod998244353」で行う。mod計算については、過去の投稿を参考にしてね。まぁ、今回はmodで複雑なことはしないけど、なぜ、「998244353」なの？という素朴な疑問を持った人には参考になると思う。
dpで解く
- 入力例は以下の通り
10 5 // 配列サイズN=10、部分配列合計NG数K=5
-5 -1 -7 6 -6 -2 -5 10 2 -10 // 配列の中身
- 回答のポイントは、「累積和配列」「dp導入」の2点。
- 部分配列の合計がKになることを防ぐのだから、工夫しないで解くと、配列要素の合計値計算が都度必要になって、すぐにTLE。よって、累積和配列を導入する。元の配列をvsとすれば、累積和配列Sは、以下の通り。
var S:[Int] = [0]
var sum = 0
for i in 0..<N {
sum += vs[i]
S.append(sum)
}
- つぎに、この問題について考えたとき、配列サイズNでの答えを求めるのではなく、配列サイズ1の場合から徐々にNに近づけられないか考える。この問題は明らかに出来そうだよね。
- 例えば、インデックスtの左側までの配列(0..<t)での解答をdp[t]として、以下の通り書ける。
//準備
var dp=[Int](repeating:0,count:N+1)
//初期値
dp[0] = 1
//遷移
for i in 1...N {
for j in 0..<i {
if S[i] - S[j] != K { // 部分配列合計がKの時は不可
dp[i] += dp[j]
dp[i] = dp[i] % 998244353 // mod計算
}
}
}
// 解答
print(dp[N])
- うむ、楽勝だね！まぁ、TLEなのだけどね。上記で提出したら、21/40でTLE。
- 問題文に付いてる入力例だけならACだったけどね！
- 制約条件は、$N≦2×10^5$であるところ、2重forループのせいで$O(N^2)$となるので、まぁ$10^{10}$を超えちゃってるよね。
計算量の削減
- さて、どうすれば、二重forループの呪縛から、逃れられるのだろうか。
- よく見ると、内側のforループは、主にif文の為だけに実施してるだけのような気が...
- all = Σdp[j] (j in 0..<i) とすれば、
- dp[i] = all - dp[?] {?:S[?] == S[i] - K}
- このようなインデックス?は存在すると限らない。S[j]{j in 0..<i}の中にちょうどS[i]-Kと一致するものがあるかどうかは運次第。
- 上記のように表現できれば、内側のforループが不要になるね。
- all = Σdp[j] (j in 0..<i)は、簡単に表現できると思う。具体的に言えば、過去のdpを全て合計しているだけだから、配列にする必要もなくて、var all = 0で初期化して、forループでdp[i]に相当する値を加算していくだけでよさそう。
- dp[?] {?:S[?] == S[i] - K} については、どうだろうか？もっと簡単に表現できないかな？上記allの表現の通り、メモ化配列について、dp[i]のようにインデックスに紐付ける必要がなくなっている。だから、インデックスiの手前までの配列に限定した場合の解答をdp[i]と表現せず、このようなイメージの辞書型配列で表現し直してみる。
- dic[S[0]:dp[0],S[1]:dp[1],S[2]:dp[2],S[3]:dp[3],...,S[i-1],dp[i-1]]
- このとき、dp[?]の代わりにdic[S[i] - K]とすれば、S[j]==S[i]-Kを満たすjがあれば、dp[j]がかえされ、そのようなjが存在しなければ、dp[?]は0が返される。
- この方針に基づいて、内側のforループをなくすようなコードを書くと、次の通り
var dic:[Int:Int] = [0:1]
var all = 1 // dicのvalueの合計だけど、いちいちdicから計算するとforループ削除した意味ないので外だし
for i in 1...N {
let ng = S[i] - K
let ans_i = all - (dic[ng] ?? 0) // swiftの辞書型はオプショナル型
dic[S[i],default:0] += ans_i
all += ans_i
if i == N {print(ans_i)}
}
- 上記コードで、dicもallも過去の解dp[j]{j in 0..<i}を全てため込んでいる。
- なお、all = dic.values.reduce(0,+)だけど、reduceで計算量を使うとTLEになってしまうので、dicと別にall単品で変数を作っている。
- .reduce(0,+)は単純な辞書のvalueの合計だけど、swiftはsumメソッドがないので、.reduce(0,+)しないといけない。sumメソッドくらい、標準で用意してくれよ。
- よし、これで提出っと...おおっと、modを忘れてたよ。
for i in 1...N {
let ng = S[i] - K
let ans_i = (all - (dic[ng] ?? 0)) % mod
dic[S[i],default:0] += ans_i ; dic[S[i],default:0] %= mod
all += ans_i ; all %= mod
if i == N {print(ans_i)}
}
-
よし！TLEがなくなった！代わりにWA(wrong answer)が発生！.....マヌケな話だけど、1時間くらい悩んで分からず諦めて、翌日さらに1時間くらい悩んで、やっと気づいたよ。
-
swiftで負数の剰余は「負数」になるってね。ちなみにc++も「負数」だけど、pythonでは「正数」となる。
swift c++ python -4 % 3 ⇒ -1 -4 % 3 ⇒ -1 -4 % 3 ⇒ 2
-
-
だから正解はこれ
for i in 1...N {
let ng = S[i] - K
let ans_i = (all - (dic[ng] ?? 0)) % mod
dic[S[i],default:0] += ans_i ; dic[S[i],default:0] %= mod
all += ans_i ; all %= mod
if i == N {print(ans_i >= 0 ? ans_i : ans_i + mod)} // 負数対応！
}
- 無事、130ms程度でAC。しかも、公式解説でのc++の解答で190msだから、swiftの方が早い！！！
おわりに
- 今回、計算量削減する為に使ったテクニックはなんて呼べば良いんだろう。
- 遷移後のdp[i]を「all - dic[ng]」で求めているから、allを全体集合と捉えて、余集合方式とか？まぁ、名付けなんて余り意味ないけど、多用するテクニックなら記憶しておきたいから、名前があった方が頭の中で整理しやすいよね。
- それにしても、5完を目指すならこの程度まで解けないといかんのね。道は遠いなぁ。
コード全体
let (N,K) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
let vs = readLine()!.split(separator:" ").map{Int($0)!}
let mod = 998244353
var S:[Int] = [0]
var sum = 0
for i in 0..<N {
sum += vs[i]
S.append(sum)
}
var dic:[Int:Int] = [0:1]
var all = 1 // dicのvalueの合計だけど、いちいちdicから計算するとforループ削除した意味ないので外だし
for i in 1...N {
let ng = S[i] - K
let ans_i = (all - (dic[ng] ?? 0)) % mod
dic[S[i],default:0] += ans_i ; dic[S[i],default:0] %= mod
all += ans_i ; all %= mod
if i == N {print(ans_i >= 0 ? ans_i : ans_i + mod)}
}