D - Handstand

競プロ 累積和

はじめに

競技プログラミング、AtCoder Beginner Contest 124
D - Handstand
言語はC#
何かあればTwitter→@pirorirori_n712まで

問題

問題文

N 人の人が左右一列に並んでいます。
0, 1 からなる長さ N の文字列 S と正整数 K が与えられます。
左から i 番目の人は、S の i 文字目が 0 のとき直立し、1 のとき逆立ちしています。
あなたは K 回まで以下の指示を行います。一度も行わなくても構いません。

指示: 
1≤l≤r≤N を満たす整数 l,r を選ぶ。
左から l,l+1,...,r 番目の人の状態を反転する。
すなわち、i=l,l+1,...,r について、左から i 番目の人は直立していれば逆立ちし、逆立ちしていれば直立する。

K 回までの指示で、逆立ちした人を連続で最大何人並ばせることができるか求めてください。

制約

・N は 1≤N≤10^5 を満たす整数である。
・K は 1≤K≤10^5 を満たす整数である。
・文字列 S の長さは N である。
・文字列 S の各文字は 0 または 1 である。

入力

入力は以下の形式で標準入力から与えられる。

N K
S

出力

K 回までの指示で、逆立ちした人を連続で最大何人並ばせることができるか出力せよ。

解き方

以下を参考にしています。
https://img.atcoder.jp/abc124/editorial.pdf
競プロ参戦記 #41 Handstand | ABC124 - Qiita
AtCoder ABC 124 D - Handstand (400 点) - けんちょんの競プロ精進記録

解き方の流れ

以下では、累積和で解いています。

問題文より、
指示を出す際に
人の状態を反転する範囲を自由に決められることが分かるので
「逆立ちした人を連続で並ばせる数」を増やすには
直立している人の連なっている範囲を指定し、
直立の人を増やさずに
逆立ちしている人のみを増やしていくことを考えます。

また、
直立を逆立ちに変えるなら
一番近い直立の範囲を反転していきたいです。

出力例2を例とすると

入力例2
14 2
11101010110011

指示を出す際に
l=9,r=12で反転させると

11101010001111

となってしまい、
逆立ちの最大の連なりは4になりますが
l=11,r=12で反転させると

11101010111111

逆立ちの最大の連なりを6にすることができます。

また、
上の指示の後の
二回目の指示を
l=4,r=4とすることもできますが

11111010111111

というように
先程反転したものとは連ならない位置に
なってしまうため
逆立ちの連続する人数を最大化するには
指示の意味がなくなってしまいます。

そのため、
反転した位置と近い位置にあるものを
反転させて連続した逆立ちの数を増やしたいです。

二回目の指示を
l=8,r=8とすれば

11101011111111

逆立ちの連なりをさらに増やし、8にすることができます。

そのため、
直立、逆立ちがそれぞれ連なっている範囲を一つの塊とみなし、
その中の直立のどの範囲を逆立ちにするかの
累積和をとっていくことで答えを出すことができます。

以下の流れで行います。

① 直立、逆立ちそれぞれの連なっている部分を一つの塊と考え、リストに保存
② 累積和をとる

① 直立、逆立ちそれぞれの連なっている部分を一つの塊と考え、リストに保存

入力の保存

var nk = Console.ReadLine().Split(' ').Select(x => Int32.Parse(x)).ToArray();
var N = nk[0];
var K = nk[1];
var S = Console.ReadLine();

リストに0と1それぞれの塊の連なりの数を保存

//直立と逆立ちの連なりの数を保存しておくリスト
var list = new List<int[]>();
//現在カウントしている数（0か1）
var currentNum = -1;
//現在カウントしている連なりのカウント
var currentNumCount = 0;
//0（直立）の連なりの数
var zeroCount = 0;
//0（直立）で始まっているかどうか
var isFirstZero = false;
//直立と逆立ちの連なりの塊を1としたときの、Sの中の塊の数
var allCount = 0;

//リストに0と1それぞれの塊の連なりの数を保存
for (int i = 0; i < N; ++i)
{
    //前の数字と今の数字の数が違ったら＝連なりが途切れたら
    if (currentNum != S[i])
    {
        allCount++;
        if (currentNum == -1) isFirstZero = S[i] == '0';
        if (S[i] == '0') zeroCount++;
        if (currentNum != -1) list.Add(new int[2] { currentNum, currentNumCount });
        currentNum = S[i];
        currentNumCount = 1;
    }
    //前の数字と今の数字の数が同じ
    else currentNumCount++;

    //最後の連なりは次に異なる数字が来ないので、indexが最後の場合に保存
    if (i == N - 1) list.Add(new int[2] { currentNum, currentNumCount });
}

変数が多いですが
後の累積和の部分で使用するためのものです。
ここで、重要なことは
0と1の連なりを一つの塊とみなし、
listに0か1かどうかと、
いくつ連なっているのかの数を保存しているということです。

② 累積和をとる

累積和をとります。

//答え
var ans = 0;
//isFirstZero のときlist内の偶数番目が0、そうでないとき奇数番目が0
//一番初め
var c = 0;
if (isFirstZero)
{
    var iC = 2 * K - 1;
    for (int i = 0; i <= iC; ++i) if (i < allCount) c += list[i][1];
}
else
{
    var iC = 2 * K;
    for (int i = 0; i <= iC; ++i) if (i < allCount) c += list[i][1];
}
if (ans < c) ans = c;

//直立を逆立ちにするパターンの数
var count = zeroCount - K + 1;
//累積和
for (int i = 1; i < count; ++i)
{
    //引く
    if (i == 1 && isFirstZero)
    {
        c -= list[0][1];
    }//偶数番目が0
    else if (isFirstZero)
    {
        c -= (list[(i - 1) * 2][1] + list[(i - 1) * 2 - 1][1]);
    }//奇数番目0
    else
    {
        c -= (list[(i - 1) * 2][1] + list[(i - 1) * 2 + 1][1]);
    }

    //足す
    //偶数番目が0
    if (isFirstZero)
    {
        var addIndex = (i + K - 1) * 2;
        c += list[addIndex][1];
        if (addIndex + 1 < allCount) c += list[addIndex + 1][1];
    }//奇数番目が0
    else
    {
        var addIndex = (i + K - 1) * 2 + 1;
        c += list[addIndex][1];
        if (addIndex + 1 < allCount) c += list[addIndex + 1][1];
    }
    //答えよりも大きかったら答えを更新
    if (ans < c) ans = c;
}

一番初めのところでは累積和の始点となる
単純に前からK個0の塊を1に変えたときの
逆立ちの連なりの数を出しています。

/答え
var ans = 0;;
//isFirstZero のときlist内の偶数番目が0、そうでないとき奇数番目が0
//一番初め
var c = 0;
if (isFirstZero)
{
    var iC = 2 * K - 1;
    for (int i = 0; i <= iC; ++i) if (i < allCount) c += list[i][1];
}
else
{
    var iC = 2 * K;
    for (int i = 0; i <= iC; ++i) if (i < allCount) c += list[i][1];
}
if (ans < c) ans = c;

listの偶数番目が0のときと
奇数番目が0のときでは計算方法が違うので分けています。

次に、
累積和をとり
前から後ろへと
0を1にする範囲をずらしています。

//直立を逆立ちにするパターンの数
var count = zeroCount - K + 1;
//累積和
for (int i = 1; i < count; ++i)
{
    //引く
    if (i == 1 && isFirstZero)
    {
        c -= list[0][1];
    }//偶数番目が0
    else if (isFirstZero)
    {
        c -= (list[(i - 1) * 2][1] + list[(i - 1) * 2 - 1][1]);
    }//奇数番目0
    else
    {
        c -= (list[(i - 1) * 2][1] + list[(i - 1) * 2 + 1][1]);
    }

    //足す
    //偶数番目が0
    if (isFirstZero)
    {
        var addIndex = (i + K - 1) * 2;
        c += list[addIndex][1];
        if (addIndex + 1 < allCount) c += list[addIndex + 1][1];
    }//奇数番目が0
    else
    {
        var addIndex = (i + K - 1) * 2 + 1;
        c += list[addIndex][1];
        if (addIndex + 1 < allCount) c += list[addIndex + 1][1];
    }
    //答えよりも大きかったら答えを更新
    if (ans < c) ans = c;
}

for文のループの回数を

//直立を逆立ちにするパターンの数
var count = zeroCount - K + 1;

としていますが、
これは例えば0の塊が5つあるとして
Kが3だとすると
5つの塊の中から連続して
3つを選ぶパターンは
5-3+1=3通りあるので
これをループをまわす回数としています。

ループの中では
K個に入る部分を入れ替えるため、
逆立ちさせる範囲に入らないものを引き
新たに逆立ちさせる範囲に入るものを足しています。

ここでも、
listの偶数番目が0のときと
奇数番目が0のときでは計算方法が違うので分けています。

後は、カウントしたものが
保存してあるansよりも大きいときに
ansの値を更新すれば答えが出ます。

全体のコード

using System;
using System.Linq;
using System.Collections.Generic;

class Dmondai
{
    static void Main()
    {
        var nk = Console.ReadLine().Split(' ').Select(x => Int32.Parse(x)).ToArray();
        var N = nk[0];
        var K = nk[1];
        var S = Console.ReadLine();
        //直立と逆立ちの連なりの数を保存しておくリスト
        var list = new List<int[]>();
        //現在カウントしている数（0か1）
        var currentNum = -1;
        //現在カウントしている連なりのカウント
        var currentNumCount = 0;
        //0（直立）の連なりの数
        var zeroCount = 0;
        //0（直立）で始まっているかどうか
        var isFirstZero = false;
        //直立と逆立ちの連なりの塊を1としたときの、Sの中の塊の数
        var allCount = 0;

        //リストに0と1それぞれの塊の連なりの数を保存
        for (int i = 0; i < N; ++i)
        {
            //前の数字と今の数字の数が違ったら＝連なりが途切れたら
            if (currentNum != S[i])
            {
                allCount++;
                if (currentNum == -1) isFirstZero = S[i] == '0';
                if (S[i] == '0') zeroCount++;
                if (currentNum != -1) list.Add(new int[2] { currentNum, currentNumCount });
                currentNum = S[i];
                currentNumCount = 1;
            }
            //前の数字と今の数字の数が同じ
            else currentNumCount++;

            //最後の連なりは次に異なる数字が来ないので、indexが最後の場合に保存
            if (i == N - 1) list.Add(new int[2] { currentNum, currentNumCount });
        }

        //答え
        var ans = 0;
        //isFirstZero のときlist内の偶数番目が0、そうでないとき奇数番目が0
        //一番初め
        var c = 0;
        if (isFirstZero)
        {
            var iC = 2 * K - 1;
            for (int i = 0; i <= iC; ++i) if (i < allCount) c += list[i][1];
        }
        else
        {
            var iC = 2 * K;
            for (int i = 0; i <= iC; ++i) if (i < allCount) c += list[i][1];
        }
        if (ans < c) ans = c;

        //直立を逆立ちにするパターンの数
        var count = zeroCount - K + 1;
        //累積和
        for (int i = 1; i < count; ++i)
        {
            //引く
            if (i == 1 && isFirstZero)
            {
                c -= list[0][1];
            }//偶数番目が0
            else if (isFirstZero)
            {
                c -= (list[(i - 1) * 2][1] + list[(i - 1) * 2 - 1][1]);
            }//奇数番目0
            else
            {
                c -= (list[(i - 1) * 2][1] + list[(i - 1) * 2 + 1][1]);
            }

            //足す
            //偶数番目が0
            if (isFirstZero)
            {
                var addIndex = (i + K - 1) * 2;
                c += list[addIndex][1];
                if (addIndex + 1 < allCount) c += list[addIndex + 1][1];
            }//奇数番目が0
            else
            {
                var addIndex = (i + K - 1) * 2 + 1;
                c += list[addIndex][1];
                if (addIndex + 1 < allCount) c += list[addIndex + 1][1];
            }

            //答えよりも大きかったら答えを更新
            if (ans < c) ans = c;
        }
        Console.WriteLine(ans);
    }
}