ここまでのあらすじ
若干手間取りながらもA問題をPythonで解き、B問題をとりあえずシミュレーションすればおkと解き、C問題で「おっやべぇグラフ系問題じゃん」と一瞬思いながらもよくよく考えれば普通にシミュレーションすればいいじゃんと気づいて解き、残り時間は30分だったか40分だったかでD問題に突入した。
問題
思考
要するに、「ある部分の総和がPの区間、総和がQの区間、総和がRの区間の3つの区間が連続して*1存在するかどうか調べろ、という問題である。
そこで、「まずはPの区間が存在するかどうかを調べ、存在するならばQの区間とRの区間について調べる」という方針をとった。
…とはいってもどうやって調べよう?
しゃくとり法
結論から言うと、しゃくとり法という考え方を使った。
var pl = 0 pr = 0 pnow = 0 while true: pnow = A[pl..pr].sum if pnow < P: pr.inc # pr += 1と同じ elif pnow > P: pl.inc else: ... pl.inc
- 調べる区間の左端、右端の位置を表す変数
pl
とpr
を用意する。最初はpl == pr == 0
である。 - 区間の総和がP未満であれば、総和が足りないから増やすために右端を1つ進める。
- 逆にPより多いときは、総和を減らすために左端を1つ進める。
- どちらでもないときはPに等しいときなので、QとRの判定に進む。判定に失敗したときは無限ループ防止のため左端を1つ進める。
こんな感じの、「左端を固定して右端を進めたり、右端を固定して左端を進めたりしながら配列を探索する」手法をしゃくとり法という*2。
制約よりなので、右端がそれを超えちゃったら終わりである。
if pnow < P: pr.inc if pr > N-3: echo "No" break
さて、あとはQとRを判定するのだが…3つの区間は連続していないといけないので、左端を進めないしゃくとり法をやろう。すなわち、オーバーしたら左端を進めるのではなくまたPの探索に戻るのである。そしてQもRも存在するとわかったらYesを出力して終わり。
... else: block checkQandR: var qnow = 0 for j in (pr+1)..(N-2): qnow += A[j] if qnow > Q: break checkQandR elif qnow < Q: continue else: var rnow = 0 for k in (j+1)..(N-1): rnow += A[k] if rnow > R: break checkQandR elif rnow < R: continue else: echo "Yes" break checkP pl.inc
というわけで完成したものがこちら
import strformat, macros, sequtils, strutils import math, algorithm, bitops, sets, deques, tables # Nim standard imput macro # example & original sourse => ttps://foo-x.com/blog/nim-contest-macro/ # arranged by stunniita let readNext = iterator(getsChar: bool = false): string {.closure.} = for s in stdin.readAll.split.filterIt(it != ""): if getsChar: for i in 0..<s.len(): yield s[i..i] else: yield s proc read(t: typedesc[string]): string = readNext() proc read(t: typedesc[char]): char = readNext(true)[0] proc read(t: typedesc[int]): int = readNext().parseInt proc read(t: typedesc[float]): float = readNext().parseFloat macro read(t: typedesc, n: varargs[int]): untyped = var repStr = "" for arg in n: repStr &= &"({arg.repr}).newSeqWith " parseExpr(&"{repStr}read({t})") macro read(ts: varargs[auto]): untyped = var tupStr = "" for t in ts: tupStr &= &"read({t.repr})," parseExpr(&"({tupStr})") macro read(n: int, ts: varargs[auto]): untyped = for typ in ts: if typ.typeKind != ntyAnything: error("Expected typedesc, got " & typ.repr, typ) parseExpr(&"({n.repr}).newSeqWith read({ts.repr})") let N,P,Q,R = read(int) var A:seq[int] = read(int,N) pl = 0 pr = 0 pnow = 0 block checkP: while true: pnow = A[pl..pr].sum if pnow < P: pr.inc if pr > N-3: echo "No" break elif pnow > P: pl.inc else: block checkQandR: var qnow = 0 for j in (pr+1)..(N-2): qnow += A[j] if qnow > Q: break checkQandR elif qnow < Q: continue else: var rnow = 0 for k in (j+1)..(N-1): rnow += A[k] if rnow > R: break checkQandR elif rnow < R: continue else: echo "Yes" break checkP pl.inc
さーて提出っと…あれ?
駄目みたいですね…仕方がないので高速化を考えます
累積和
pnow = A[pl..pr].sum
よくよく考えたらこうやっていちいち区間の総和を求めていたらクッソ時間がかかってしまいますね…そこであらかじめ累積和を取っておきましょう。
var B = newSeq[int](N+1) B[1] = A[0] for i in 1..<N: B[i+1] = B[i] + A[i]
入力例1だとこんな感じに*3
こうしておくことで、
pnow = B[pr+1] - B[pl]
というふうに、区間の和を一発で求められます。やったね。凄いね。
ってことで提出…勝ったなガハハ
それがこのざまである。
まさかの1個だけTLE…いやどうすんねんこれ以上改良する箇所なんて…
…いや、心当たりはあります。QとRの判定の部分です。
var qnow = 0 for j in (pr+1)..(N-2): qnow += A[j] if qnow > Q: break checkQandR elif qnow < Q: continue else: ...
これ…結局のところやっているのは線形探索なので…。たぶん、いや間違いなくここが原因でしょう。
でもどうすりゃいいんだ?…結局時間内には思いつくことは無かった…。
二分探索
「累積和を計算してあるんだからそれを利用して二分探索すれば間に合うんじゃね?」
だが時既に時間切れ。仮にコンテスト中に思いついたとしても実装は間に合ってなかったであろう。
二文探索には二種類ある*4。
- 二分探索(原義)…ソート済み配列からある値が存在するかどうか、あるいはある値がどこに入るかを高速で調べる
- 二分探索(広義)…ある条件を満たすかどうかがある値を境目に二分されているとき、ソート済み配列において、その境界線を高速で調べる
前者は単純でありライブラリが既に実装済みだったりするが、後者は自力で実装しないといけない面倒なものであり、ここで必要なのは後者だと思われたからである。区間和がQ(R)より大きいか小さいかなどという条件を調べるのはライブラリでは無理じゃね?
結局すぐには解決できず、この後実生活が忙しくなったこともあり、しばらくこの問題からは離れることとなった。
後日
忙しい間もこの問題のことが頭の片隅でずっと引っかかっていた…そしてある日気が付いた
「いやこれ原義二分探索でもいけるんじゃね?」
もう一度入力例1を使って考えてみよう。
Pを探索していくとpl = 1, pr = 2
においてpnow = B[pr+1] - B[pl] = 6 - 1 = 5
となる。この後も調べていくとB[6] - B[3] = 13 - 6 = 7
であり、B[8] - B[6] = 18 - 13 = 5
であり、答えはYesである。ここからがポイントで、「Pの隣にQの区間があることと、Pの右端の値にQを足したものがBに存在することは、同じ」なのだ。13 - 6 = 7と6 + 7 = 13は同じだし、13はBに存在する。
つまり、B[pr+1] + Q
がBに存在するかどうかを調べればよいのだ。そうと決まれば話は早い。nimのbinarySearch
はその値が配列に存在すればその添え字を、存在しなければ-1
を返すので、B.binarySearch(B[pr+1] + Q) != -1
とすればよい。Rについても同様にB.binarySearch(B[pr+1] + Q + R) != -1
である。
というわけで出来上がったものがこちら
# ここはさっきと同じなので省略 let N,P,Q,R = read(int) var A:seq[int] = read(int,N) pl = 0 pr = 0 pnow = 0 B = newSeq[int](N+1) #型、長さ B[1] = A[0] for i in 1..<N: B[i+1] = B[i] + A[i] #echo $B block checkP: while true: #pnow = A[pl..pr].sum pnow = B[pr+1] - B[pl] if pnow < P: pr.inc if pr > N-3: echo "No" break checkP elif pnow > P: pl.inc else: if B.binarySearch(B[pr+1] + Q) != -1 and B.binarySearch(B[pr+1] + Q + R) != -1: echo "Yes" break checkP pl.inc
あーすっきりした
おまけ
「Pも二分探索すれば良くね?」…実際、解答はこの方針を採用している。Pの左端を全探索し、そのそれぞれで二分探索を行っている。
というわけでやってみた
# ここはさっきと同じなので省略 let N,P,Q,R = read(int) var A:seq[int] = read(int,N) pl = 0 pr = 0 pnow = 0 B = newSeq[int](N+1) #型、長さ ans = false B[1] = A[0] for i in 1..<N: B[i+1] = B[i] + A[i] #echo $B for x in 0..(N-3): if B.binarySearch(B[x] + P) != -1 and B.binarySearch(B[x] + P + Q) != -1 and B.binarySearch(B[x] + P + Q + R) != -1: ans = true break if ans: echo "Yes" else: echo "No"
たぶんこれが一番楽だと思います