[AtCoder] ABC 124 D – Handstand

問題

方針

最終的な形を考える

\( K \) 回までの指示をした時の文字列において、最大何個の 1 が連続して並んでいるかを問われているので、指示の順番を考慮する必要がないことが分かります。つまり、最終的な形がどのようになるかを考えればよいです。

文字列の圧縮

“0001011" という文字列を連続している文字の個数で表現すると、 \( (3, 1, 1, 2) \) という配列が得られます。また、考えやすいように、"10011″ というような文字列の先頭が1 から始まるものは、ダミーの値として \( (0, 1, 2, 2) \) という配列で表現するようにします。このような配列を \( v \) とします。

\( v \) を \( 0 \) オリジンの配列とすると偶数番目の要素は、"0″ が連続している個数であり、奇数番目は “1" が連続している個数となります。

累積和

\( v \) の累積和を \( b \) とします。圧縮された文字列において、"0″ を “o"、"1" を “x" と表現すると、"oxoxox" のように “o" から始まる “ox" の文字列となります。ここで、次のように指示を与えることを考えます。

  • \( i \) 番目の “o" を起点として、右に向かって 最大 \( K \) 個の “o" を “x" にする

例えば、\( 0 \) 番目の “o" となっている区間に対して指示を与えると、"xxoxox" となります。全ての起点となる \( i \) に対して、\( i \) を起点として \( 1 \) が連続している個数が分かれば良いです。これは累積和を利用することで計算することができます。

起点となる \( i \) は、\( b \) の偶数番目の要素に対応します。最終的な形の圧縮された文字列の右端の要素番号 \( r \) は \( r = i + 2K \) となります。また、 \( r \) は \( b \) のサイズを超えない値とします。よって、\( i \) を起点したときの連続する “1" の個数は、

\[ b_{r} – b_{i – 1}\]

となります。ここで、\( i =0 \) のときは、\( b_{-1} = 0 \) とします。

コード

提出したコード

累積和の計算

vector<int> v;
  if(S[0] == '1') {
    v.push_back(0);
  }
  char c = S[0];
  for (int i = 0; i < N; i++) {
    int cnt = 1;
    while(i + 1 < N) {
      if(S[i + 1] == c) {
        i++;
        cnt++;
      } else {
        break;
      }
    }
    v.push_back(cnt);
    c = (c == '0') ? '1' : '0';
  }
  
  int l = v.size();
  int *b = new int[l + 1]{};
  for (int i = 0; i < l; i++) {
    b[i + 1] = b[i] + v[i]; 
  }

動的配列を使うと良いかもしれません。

最大値の更新

  int ans = 0;
  for (int i = 0; i < l; i += 2) {
    if (i == 0) {
      int r = min(l, 2 * K);
      ans = b[r];
    } else {
      int r = min(l, i + 2 * K);
      ans = max(ans, b[r] - b[i - 1]);
    }
  }

\( i = 0 \) と \( r \) が \( b \) のサイズを超えないことに注意します。