すゅん/Суюнのメモ帳/裏紙

またの名をゴミ捨て場

【解説もどき】ABC265Dを解いてみたかったけど解けなかったので後日解いた

ここまでのあらすじ

若干手間取りながらもA問題をPythonで解き、B問題をとりあえずシミュレーションすればおkと解き、C問題で「おっやべぇグラフ系問題じゃん」と一瞬思いながらもよくよく考えれば普通にシミュレーションすればいいじゃんと気づいて解き、残り時間は30分だったか40分だったかでD問題に突入した。

問題

atcoder.jp

思考

要するに、「ある部分の総和が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
  1. 調べる区間の左端、右端の位置を表す変数plprを用意する。最初はpl == pr == 0である。
  2. 区間の総和がP未満であれば、総和が足りないから増やすために右端を1つ進める。
  3. 逆にPより多いときは、総和を減らすために左端を1つ進める。
  4. どちらでもないときはPに等しいときなので、QとRの判定に進む。判定に失敗したときは無限ループ防止のため左端を1つ進める。

こんな感じの、「左端を固定して右端を進めたり、右端を固定して左端を進めたりしながら配列を探索する」手法をしゃくとり法という*2

制約より x \leq N-3なので、右端がそれを超えちゃったら終わりである。

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

さーて提出っと…あれ?

TLE

駄目みたいですね…仕方がないので高速化を考えます

累積和

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

\displaystyle{
A = [1, 3, 2, 2, 2, 3, 1, 4, 3, 2 ]  \rightarrow B = [0,1,4,6,8,10,13,14,18,21,23 ]
}

こうしておくことで、

pnow = B[pr+1] - B[pl]

というふうに、区間の和を一発で求められます。やったね。凄いね。

ってことで提出…勝ったなガハハ

TLE

それがこのざまである。

まさかの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を使って考えてみよう。

\displaystyle{
A = [1, 3, 2, 2, 2, 3, 1, 4, 3, 2 ]  \rightarrow B = [0,1,4,6,8,10,13,14,18,21,23 ]
}

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"

たぶんこれが一番楽だと思います


*1:連続しているかどうかを考えない嘘解法があったらしいが、コンテスト後にテストケースが追加され、現在ではWAになる…らしい。

*2:しゃくとり虫が進む様子に例えられてこの名前がついた…らしい。ちなみに、右端も左端も最終的に配列の端から端まで移動するので計算量は O(N)

*3:Bの最初に0を仕込んでおいたのはなんかうまくいかなかったからである

*4:この原義と広義の分け方は自分が勝手に呼んでいるだけなので他人には通じない可能性が高い。