【図解】最難関Sランクの5問目「十億連勝」を解いてみた。O(N)解法あり。

はじめに

paiza×Qiita記事投稿キャンペーン「プログラミング問題をやってみて書いたコードを投稿しよう！」
というのをやってます。参加してみました。

やった問題

• ステージがいくつかある
• それぞれのステージに試合がいくつかある
• ステージ内の試合の順番はあらかじめ決められている
• ステージ1から順に挑戦する
• 負けたら、そのステージに残り試合があっても次のステージに移動
• 全ステージが終わったときに連勝の最高がX連勝なのは何通りか

結果

100点でした。

解説

ちょっと入力が特殊なのでまずそれに対応する。

ステージに試合がない場合

0 ≦ A_i ≦ 1,000,000,000 という条件なのでステージに試合がない場合がある。
これは単純にそのステージはないことにする。

隠し要素の条件が0の場合

0 ≦ X ≦ 1,000,000,000 なので隠し要素の条件が0の場合がある。
そのときはすべての試合に負けること1通りなので1を出力。
また、試合自体が全くない場合でも「N 個すべてのステージを終えた時点で、それまでの最大の連勝数がちょうど X だった場合」は1通りである。

試合がひとつもない場合

隠し要素の条件が0の場合は1通りだが
その他の場合は0通り。
実は提出コードでミスったが不正解にならなかったのでこれは出題者の意図ではないのかもしれない。

以上の特殊な状況を早期に省くと

• ステージに1個以上試合がある
• ステージは1つ以上ある
• 隠し要素の条件は1以上

になる。

---

基本方針

連勝が X 以下のものA個
連勝が X-1 以下のものB個
だとA-Bが答えになる。
この方針にする。

ステージカウンターを導入する。これはあるステージからゲームを始めたときに(前のステージのことは考慮しない)、条件を満たすものがいくつあるか、というもの。条件とは「連勝がxxx以下」のこと。この値を最終ステージから逆向きに測定していき、やがて初めのステージに欲しい値が設定される。

計算のため、ステージカウンターは最後のステージの後ろに1というデータを付け足す。

ステージカウンター(sc)に設定する処理

func count(xxx: Int, s: Int) -> Int

この関数は、ステージ s 以降を対象にして条件にあうのが何通りか、を返す。与えられる xxx により、そのステージ s の頭からの連勝を xxx 以下に制限する。この xxx は入力(関数の入力ではなくプログラムの標準入力)で与えられる x とは別の値になることがある点に注意する。xxx が7のとき、ステージ s の頭から4連勝(5試合目は負け)して、ステージs+1以降でそれと矛盾しないようなものが100個あれば、4連勝の分を100個とカウントする。

この関数の中身を説明する。

基本的に「どこに負けを置くか」という考え方をする。
xxx が3なら0勝、1勝、2連勝、3連勝の後ろに負けを置く。つまり負けの置き方は4つ。で、負けを置く場所がそのステージだけで終わるか次のステージに行くかで場合分けをする。

そのステージに負けを置くだけで計算が終わる場合

xxx が3で、そのステージの試合が5なら、そのステージに負けを置く。
ここになるには、xxx がステージの試合数より小さいとき(同じだと負けを置けないことがある)。
ある負けの置き方に対し、矛盾しないステージs+1以降の置き方はステージカウンターの sc[s + 1] 個。
負けの置き方が xxx + 1 通りなので、それに sc[s + 1] をかけた (xxx + 1) * sc[s + 1] が求める数。すでに計算してあるステージカウンターを利用できる。

負けの置き方が次のステージにも延期される場合

xxx が7で、そのステージの試合が5なら、次のステージに引き継ぐことを考えないといけない。
まず、xxx が7でも、3連勝して4つ目に負けるというのもカウントする。このようにこのステージに負けを置く場所がそのステージの試合数と同じだけある。これのそれぞれにステージs+1のステージカウンターをかけた 試合数 * sc[s + 1] をメモる。
そして次のステージに渡すため count(xxx: Int, s: Int) を再度を呼び出し count(xxx: xxx-試合数, s: s + 1) とする。この戻り値をさっきのメモに足す。以上より sc[s] が求まる。

次へ渡すのが0のこともある。

ただし

次のステージに渡そうとしても今のステージが最後のステージの場合は、次に投げることができない。この場合は最後のステージの最後の試合に勝った1通りが加算される。

コード

Swiftです。

import Foundation

let mod = 1_000_000_000

func readInt() -> Int {
    Int(readLine()!)!
}

func readInts() -> [Int] {
    readLine()!.split(separator: " ").map { Int(String($0))! }
}

let nx = readInts()
var n = nx[0]
let x = nx[1]

var a: [Int] = []
for _ in 1...n {
    let s = readInt()
    if s != 0 {
        a.append(s)
    }
}
n = a.count

if x == 0 {
    print(1)
    exit(0)
}

if n == 0 {
    print(0)
    exit(0)
}

func count(xx: Int) -> Int {
    var sc = [Int](repeating: 0, count: n + 1) //stage counter
    sc[n] = 1
    
    func count(xxx: Int, s: Int) -> Int {
        //ステージsの先頭からxxx以内の連勝するのが何通りか
        
        var returnVal = 0
        
        //どこに負けを置くか
        
        if xxx < a[s] {
            //このステージに負けを置いて終了する
            returnVal += (xxx + 1) * sc[s + 1]
        } else {
            
            //このステージに負けを置く
            returnVal += a[s] * sc[s + 1]
            
            //次のステージ以降に負けを置く
            if s + 1 < n {
                //次のステージに渡す
                returnVal += count(xxx: xxx - a[s], s: s + 1) 
            } else {
                //最後のステージの最後のゲームに勝って終わり
                returnVal += 1
            }
        }
        returnVal %= mod
        return returnVal
    }
    
    for s in (0..<n).reversed() {
        sc[s] = count(xxx: xx, s: s)
    }
    
    return sc[0]
}

let valueX1 = count(xx: x)
let valueX2 = count(xx: x - 1)

print((valueX1 + mod - valueX2) % mod)

感想

もっといいのがあるかもしれないが、N=4000なので O(N^2) でとりあえず提出した。

追記 O(N)の解法

paiza Qiita コラボキャンペーンの終了後に追記しております。要は締め切りの後で追記しています。

上で紹介しているO(N^2)解法を改良し、O(N)の解法が出来ました。同じ計算を繰り返しているところを解消しました。

O(N^2)の解法は、次の図のように負けを置くことができる試合に負けを置いて、それぞれの値(次のステージのステージカウンターsc)の合計を出すものでした。

上の図は、ステージ2のステージカウンター(sc)を調べているところです。負けを置ける場所は10個です。

ステージ3を調べる状態も並べてみます。

この2つは共通しているところがあります。下の図の青で塗ったところが共通してます。

この何度も同じことをやる計算を1回にするためにステージサムカウンターsscを作ります。このsscもscと同様に各ステージごとに持ちます。
ある場所に負けを置いた値は常に同じです。その値を最後からそのステージまで足したものがsscです。

そのステージの試合の数と次のステージのscをかけたものをそのステージのsscに加算します。後ろのステージのsscも加算します。

決定する順番は次のようになります

• ステージ5のssc
• ステージ5のsc
• ステージ4のssc
• ステージ4のsc
• ステージ3のssc
• ステージ3のsc
• ステージ2のssc
• ステージ2のsc
• ステージ1のssc
• ステージ1のsc

scの求め方は次のようになります。

上の図はステージ2のステージカウンターsc1を求めています。

ステージ2以降に負けを置いた値をすべて足した値ssc1から
ステージ4以降に負けを置いた値をすべて足した値ssc3を引きます。
これで最初の6個の負けの分が求まります。
ステージ4の2個の負けは中途半端になっているのでそこは計算します。これは2個かけるsc4です。これを足します。
sc1が求まりました。

import Foundation

let mod = 1_000_000_000

func readInt() -> Int {
    Int(readLine()!)!
}

func readInts() -> [Int] {
    readLine()!.split(separator: " ").map { Int(String($0))! }
}

let nx = readInts()
var n = nx[0]
let x = nx[1]

var a: [Int] = []
for _ in 1...n {
    let s = readInt()
    if s != 0 {
        a.append(s)
    }
}
n = a.count

if x == 0 {
    print(1)
    exit(0)
}

if n == 0 {
    print(0)
    exit(0)
}

var start = [Int](repeating: 0, count: n)
var end = [Int](repeating: 0, count: n)

for s in 0..<n {
    if s > 0 {
        start[s] += end[s - 1] + 1
    }
    end[s] += start[s] + a[s] - 1
}

func count(xx: Int) -> Int {
    var sc = [Int](repeating: 0, count: n + 1) //stage counter
    sc[n] = 1
    
    //scを後ろのステージから累積したもの
    var ssc = [Int](repeating: 0, count: n + 1) //stage sum counter
    ssc[n] = 1
    
    var loseStage = n - 1 //負け試合の最大はどこのステージか
    for s in (0..<n).reversed() {
        ssc[s] = a[s] * sc[s + 1]
        ssc[s] += ssc[s + 1]
        ssc[s] %= mod
        
        let loseMatch = start[s] + xx //負け試合の最大
        if loseMatch > end[n - 1] {
            sc[s] = ssc[s]
        } else {
            while loseMatch < start[loseStage] {
                loseStage -= 1
            }
            sc[s] = (ssc[s] + mod - ssc[loseStage] + (loseMatch - start[loseStage] + 1) * sc[loseStage + 1]) % mod
        }
    }
    
    return sc[0]
}

let valueX1 = count(xx: x)
let valueX2 = count(xx: x - 1)

print((valueX1 + mod - valueX2) % mod)