ABC 124-D Handstand を累積和を用いて解いてみる(Go)

irisruneです。新しめのABC-D問題を扱うのは久々ですね。

atcoder.jp

解説だけでも3種類の解き方が示されていますが、 今回用いた解き方は累積和を用いたしゃくとり法に少し近い解き方です。見た瞬間はUnion-Find木使うのかと思いました。

package main

import (
    "fmt"
)

func maxInt(a, b int) int {
    switch {
    case a >= b:
        return a
    default:
        return b
    }
}

func main() {
    var n, k int
    var str string
    fmt.Scan(&n, &k, &str)

    var partLength []int
    prevChar := str[0:1]
    cnt := 1
    for i := 1; i < n; i++ {
        if str[i:i+1] == prevChar {
            cnt++
        } else {
            partLength = append(partLength, cnt)
            prevChar = str[i : i+1]
            cnt = 1
        }
    }
    partLength = append(partLength, cnt)

    var cumSum []int
    cumSum = append(cumSum, 0)
    for i := range partLength {
        cumSum = append(cumSum, cumSum[i]+partLength[i])
    }
    if str[0:1] == "0" {
        cumSum = append([]int{0}, cumSum...)
    }
    if str[n-1:n] == "0" {
        cumSum = append(cumSum, n)
    }

    if len(cumSum) <= 2*k+2 {
        fmt.Println(n)
    } else {
        maxRet := 0
        for i := 0; i < len(cumSum)-(2*k+1); i += 2 {
            maxRet = maxInt(maxRet, cumSum[i+(2*k+1)]-cumSum[i])
        }
        fmt.Println(maxRet)
    }

}

解き方の方針は、

1. 0と1の連続する部分列の長さの配列を作る。
2. その配列について累積和の配列を作る。
3. 0が0個連続する部分列から始まり、1が任意の個数連続する部分列で終わるように累積和の配列の先頭と末尾に要素を足す。
4. について、番目と番目の累積和の差の最大値を求めて解とする。ただし、累積和の配列長が以下の場合はを解とする。

と少しややこしいものになっています。重要な処理は太字にした3.の部分ですが、4.の処理を単純にするためのものなので全体的に簡単に説明したいと思います。

入力例2について2.まで処理が終わった状況を例示してみます。 1行目は部分列の長さの配列、2行目は補足としてそれぞれ0と1とどちらの連続する部分列かを表したもの、1行空けて4行目は累積和の配列です。

3 1 1 1 1 1 2 2 2
1 0 1 0 1 0 1 0 1

3 4 5 6 7 8 10 12 14

さて、入力例2ではなので、0の連続する部分列を1が連続する部分列に2つ変換することができます。 その場合最も1が連続する部分列が長くなるのは、1が連続する部分列で始まり1が連続する部分列で終わる個の部分列を繋げるように変換した場合です。

その部分列の長さは、末尾側の1が連続する部分列までの累積和から、先頭側の1が連続する部分列の手前の0が連続する部分列の累積和を引いたものになります。 これを求めるために、先頭に0が連続する長さ0の部分列を追加すると、

0 3 4 5 6 7 8 10 12 14

となるので、累積和の差を先頭から2個ずつインデックスをずらして求めた最大値が解となります。まあ、累積和の求め方の関係で先頭に0はどうしてもついてくるのですが。

入力例1については、元の配列の先頭が0、末尾が0のために先頭と末尾それぞれに1が連続する長さ0の部分列を追加する必要があります。さらに0が連続する長さ0の部分列が先頭に来るように、これも先頭に追加することで解を求めることができます。

入力例3については、各種処理を行った結果得られる累積和の配列の長さが2しかないので、元の配列の長さである1をそのまま解とすることになります。

以上のように各入力例についての解き方を示すことができ、残りの入力に対しても解くことができたわけですが…説明が非常に難しいですね。

コードレビュー

• 累積和周りの処理アルゴリズムがコードからわかりにくい。
• 部分列に分割するところは仕方ないとしても累積和は部分列の長さから配列長をある程度は予測できるので、makeで初期化した方が処理時間は短く済む可能性がある。

雑記

• int用のmax関数を自前で実装したり、switch文やrange文を使ったりすると結構読みやすくなってきたと思います。

• 次回はsquare869120Contestで解けた問題から一番配点（部分点）の高いものを扱う予定です。

  • パッと見て300点は最低でも解けるはず。