いもす法とニブタンのマリアージュ【競技プログラミング】

はじめに

• solve関数を使ったニブタンについて、最近投稿したけど、ニブタンがいくら高速でも、solve関数の計算が遅いとどうしようもない。
• このsolve関数をいもす法で高速化しないと解答できない問題があったので、対応したいと思う。

どんな問題？

• 第五回 アルゴリズム実技検定 過去問のm問題(棒の出荷)となる。
  • 木の棒に対して、「N-1個の切れ込み」があり、好きな位置で切ることが出来るが、
  • 切断後の棒は、それぞれ長さ「L以下」にしないといけない。
  • 切断後の棒の中で、最も短い棒の長さを最大限長くした場合、その長さはいくらか？
  • ちなみに、N-1個の切れ込み全て切断ときはN本の棒となり、その各棒の長さが､入力値として与えられる。
• 例えば、入力例は以下の通り、

4 5 // 全ての切れ込みで切ったとき4本、切断後の棒が全て5以下になるようにする
4 1 5 5 // 切れ込みを全て切断したときの棒の長さは、左から4,1,5,5

• 上記切れ込みを図示すると以下のようになる。
  

ニブタンで解く

• さて、solve関数を作ってニブタンすることは既定路線だけど、どうやって、solve関数を作ろうか？
• まぁ、最初にsolve関数を使ったニブタンを紹介したときも、同じような問題を解いていたし、同様とするね。つまり、solve関数の引数を「x」とした場合、「切断後の各棒の長さをx以上とできるか？」という条件とする。
• さて、solve関数の作り方だけど、切断後の棒の長さは、「x以上」とすると同時に「L以下」としなければならない。
  • 例えば、上図の切れ込みに対して、x=1のとき、黄(4)⇒緑(1)⇒青(5)⇒赤(5)は、全て「1以上5以下」を満たすから、切り取れる
  • x=2のときは、最初の切れ込み(オレンジ:4)は「2以上5以下」を満たすから切り取れるが、次の切れ込み(緑:1)は短すぎて切り取れない。じゃあ、その次の切れ込み(緑+水色:6)は長すぎて切れ取れない。
  • あれ？ということは、solve関数はfalseとなっちゃう？...んなわきゃない。
  • そもそも、最初のオレンジ部分で切り取らず、「オレンジ＋緑」(長さ5)で切り取ればOK。これも「1以上5以下」を満たすからね。ここで切断すれば、次の水色(5)もその次の赤色(5)も切り取れる。
• それでは、どうやって、solve関数を実装すれば良いだろうか？
  • 各切れ込みiに、「1以上5以下」のルールを守って、切り取り可能か？というtrue/falseを格納する配列check[i](i in 0...N)をつくることとする。
  • i=0は棒の左端、i=Nは棒の右端とする。初期条件はcheck[0]=trueとなり、最終的に、切り取り可能かどうかは、check[N]がtrue/falseで判定する。今回、配列名をcheckとしたけど、考え方的にはdpとしても良かったね。
• 上記でまとめた作成方針に従ってsolve関数を実装すると､以下の通り。あ、ちなみに、切れ込み毎の長さについて、累積和の配列sumsを作ったよ。

let (N,L) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
let vs = readLine()!.split(separator:" ").map{Int($0)!}

// vsの累積和
var sums:[Int] = [0]
var sum = 0
for v in vs {
    sum += v
    sums.append(sum)
}

// 全ての棒をx以上L以下の長さで切り取れるか？
func solve(_ x:Int) -> Bool {
    // dpの準備
    var check = [Bool](repeating:false,count:N+1)

    // 初期条件
    check[0] = true

    // 遷移：「切れ込みi」から「切れ込みj」の長さが「x以上L以下」なら、check[j]をtrueにする。
    for i in 0..<N { 
    if check[i] == false {continue}
        for j in i..<N {
            let t = sums[j+1] - sums[i]
            if t >= x && t <= L {
                check[j+1] = true
            }
        }
    }
    return check[N]
}

• いやぁ、dpを覚えたお陰で、理解の幅が広がったよ。dp様、助かるっす。
• で、上記のsolve関数を使ったニブタン部分は、以下の通り。

let min = 1 // 検索範囲の最小値
let max = L // 検索範囲の最大値

var ok = min - 1 // 左側の開始位置 -- ループ毎に「同位置」or「中間値」へ移動
var ng = max + 1 // 右側の開始位置 -- ループ毎に「同位置」or「中間値」へ移動

while abs(ok - ng) > 1 {
    let mid = (ok + ng) / 2
    print(mid,ok,ng,solve(mid)) // 出力を参照
    if solve(mid){
        ok = mid
    } else {
        ng = mid
    }
}

print("答え",ok) // 出力を参照

///出力
3 0 6 true
4 3 6 true
5 4 6 true
答え 5

• よし！完成！ ... あれ？いもす法は？
• 上記のsolve関数は、二重forループを使っているので、計算量は$O(N^2)$となる。んで、コンテストにおけるNの制約は、$N ≦ 2×10^5$ なので、上記のままだと、計算量は$10^8$を大幅に超えるのでTLEだね。実際に提出したら、AC:15、TLE:19でアウトでした。
• ここで、いもす法の出番です。
• 二重forループを避けたいときは？いもす法！！　だよね。（今日、いもす法を知ったばかりのくせに...調子に乗りました。ゴメンナサイ）

いもす法をsolve関数で使う

• さて、高速化を目的に、二重forループを分離するため、いもす法を導入する。
  • そもそも、いもす法は、
    • 「対象となる配列」に対して、
    • 旧外側のループで、増減のみ記録させ
    • 旧内側のループで、類型和もどきの手法で、元々の配列にする
  • はず、なんだけど....「対象となる配列」がBool型じゃ駄目じゃね？
  • よって、配列をInt型にして、0をfalse,0超をtrueの意味にするよ！
  • ところで、初歩的ないもす法では、
    • 「対象となる配列」を、最初のforループ(N個の位置情報登録)で「増減のみ記録」させ、
    • 次のforループ(全区間で、位置情報を累計化)で累計和を再構築、と言うことをしていた。
  • 今回は、若干複雑であり、
    • 「増減のみ記録」する配列と、
    • この「増減のみ記録」された配列の累計和を取る元々の配列（check配列）
      を別々にする。
  • これは、いもす法導入前のsolve関数を見てもらえば、分かり易いかもしれないけど、dpの遷移があるアルゴリズムなので、「増減のみ記録」の配列を単品で生成出来ないためである。
    • dpの遷移、とは、具体的に言えば、下図のcheck[i]==trueのとき、一定の条件(緑枠内)を満たすj+1に対して、check[j+1]=trueを伝播させる、ということ。
      
    • また、「増減のみ記録」の配列を単品で生成出来ない、とは、具体的に言えば、
    • 上図のforループのiを進める都度、累計和もどきの結果となる配列checkを生成して、
    • check[i]がtrue/falseを判定（上図の紫枠）した後に、「増減のみ記録」、つまり、
    • 上図の緑枠内の「j+1」について、tがxに等しくなるときに増加「+1」し、
    • tがLと等しくなるときに減少「-1」させる必要がある
      と言うこと。
• 以上を踏まえて、いもす法を導入したsolve関数は以下の通り。

// 全ての棒をx以上L以下の長さで切り取れるか？
func solve(_ x:Int) -> Bool {
    var check = [Int](repeating:0,count:N+1)
    var cd = [Int](repeating:0,count:N+1) // checkの差分dの意味でcd

    // 元のsolve関数の t=vs[j+1]-vs[i]を考える
    var l = 0 // t=xとなるインデックス-1を格納する変数（切断可能範囲の左端）
    var r = 0 // t=Lとなるインデックスを格納する変数（切断可能範囲の右端）

    //初期化:元のsolve関数のcheck[0]=trueに相当
    cd[0] += 1
    cd[1] -= 1

    for i in 0..<N {
    
        if i>0 { // [i-1]インデックスが有るため、対応
            check[i] = check[i-1] + cd[i] // cdのiまでの累積和Σcd[t] (t in 0...i)と同じ
        } else {
            check[i] = cd[i] // i==0の時のみ
        }
            
        if check[i] == 0 {continue} // 元のsolve関数の「check[i] == false」条件と同じ
        
        while sums[l] - sums[i] < x {l += 1} // 元のsolve関数のt=xになる直前までlを増加
        cd[l] += 1
        
        while sums[r+1] - sums[i] <= L {r += 1} // 元のsolve関数のt=Lになる直前までr+1を増加
        cd[r+1] -= 1 
    }    
    
    check[N] = check[N-1] + cd[N]
    return check[N]>0
}

• いもす法導入のための主な変更点
  • 前述の通り、「対象となる配列」(check)の中身をBoolからIntへ
  • 前述の通り、「対象となる配列」と別に、「増減のみを記録する配列」(cd)を用意する
  • 配列cdへの増減記録タイミングを計るための変数lとrを導入する
    • lとrは、0からスタートして、Nに向けて進んでいくが、これは、whileループによる。
    • whileループの条件節の意味は次のとおり。
      • lの意味：条件節より、「sums[l]-sums[i]>=x」を満たす、最小のl。
      • r+1の意味：条件節より、「sums[r+1]-sums[i]>L」を満たす、最小の「r+1」。つまり、rについては、「sums[r]-sums[i]<=L」を満たしている
    • 上記により、条件を満たすインデックスlにおいて、増加「cd[l]+=1」を記録し、条件を満たさなくなるr+1において、減少「cd[r+1]」を記録する。
• いもす法導入により、solveの計算量は、forループのN回転と、l,rの0⇒N遷移からO(N + N)になるから、結局O(N)となる。いもす法導入前の$O(N^2)$から大幅に減少できたね。
• コード全体は、以下の通り。

let (N,L) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
let vs = readLine()!.split(separator:" ").map{Int($0)!}

// vsの累積和
var sums:[Int] = [0]
var sum = 0
for v in vs {
    sum += v
    sums.append(sum)
}

// 全ての棒をx以上L以下の長さで切り取れるか？
func solve(_ x:Int) -> Bool {
    var check = [Int](repeating:0,count:N+1)
    var cd = [Int](repeating:0,count:N+1) // checkの差分dの意味でcd

    // 元のsolve関数の t=vs[j+1]-vs[i]を考える
    var l = 0 // t=xとなるインデックス-1を格納する変数（切断可能範囲の左端）
    var r = 0 // t=Lとなるインデックスを格納する変数（切断可能範囲の右端）

    //初期化:元のsolve関数のcheck[0]=trueに相当
    cd[0] += 1
    cd[1] -= 1

    for i in 0..<N {
    
        if i>0 { // [i-1]インデックスが有るため、対応
            check[i] = check[i-1] + cd[i] // cdのiまでの累積和Σcd[t] (t in 0...i)と同じ
        } else {
            check[i] = cd[i] // i==0の時だけ
        }
            
        if check[i] == 0 {continue} // 元のsolve関数の「check[i] == false」条件と同じ
        
        while l <= N && sums[l] - sums[i] < x {l += 1} // 元のsolve関数のt=xになる直前までlを増加
        if l <= N {cd[l] += 1}
        
        while r+1 <= N && sums[r+1] - sums[i] <= L {r += 1} // 元のsolve関数のt=Lになる直前までr+1を増加
        if r+1 <= N {cd[r+1] -= 1} 
    }    
    
    check[N] = check[N-1] + cd[N]
    return check[N]>0
}

let min = 1 // 検索範囲の最小値
let max = L // 検索範囲の最大値

var ok = min - 1 // 左側の開始位置 -- ループ毎に「同位置」or「中間値」へ移動
var ng = max + 1 // 右側の開始位置 -- ループ毎に「同位置」or「中間値」へ移動

while abs(ok - ng) > 1 {
    let mid = (ok + ng) / 2
    // print(mid,ok,ng,solve(mid)) // 出力を参照
    if solve(mid){
        ok = mid
    } else {
        ng = mid
    }
}

// print("答え",ok) // 出力を参照
print(ok) // 出力を参照

• 上記では、l,rの条件修正している。
  • 実は、というか当然だけど、lもr+1も「<=N」条件を満たさないとインデックス範囲外エラーになるよね。
• AtCoderに提出したら、200ms以内でACです。やったね！

おわりに

• 問題の解き方は理解できるようになったけど、コンテスト中にこんがらがらずに回答出来るかはまだ微妙かな。
• いもす法導入前のは、自力で書けると思うけど、今回のいもす法は、下記のインデックスをこんがらがらずに書ける気がしない。
  • check[i] = check[i-1] + cd[i] とか、
  • while sums[l] - sums[i] < x {l += 1}　とか、
  • while sums[r+1] - sums[i] <= L {r += 1} 　とかね
• i-1とか、r+1とか、<xとか、<=Lとかさぁ、-1,+1,<,<=を実際に全てミスなく設定できるか自信ないわ。
• 昨日のABC368のc問題も、微妙な設定ズレが解消できなくて、a,b問題を15分で終わった後、この問題だけで残り85分を使い切ってもズレを解消できず。コンテスト終了後10分後にやっと気付いて、ACとれた。
• 結局は、自分がミスしやすいところって、傾向があるから、場数を踏んで対応能力を上げるしかないのかな。