Problem p03074 (C++)

Buggy submission (status: Runtime Error):


#include <algorithm>
#include <fstream>
#include <functional>
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

#define rep(i, start, end) for (int i = (int)(start); i < (int)(end); i++)
#define charToInt(c) c - 'a'

int main(int argc, const char *argv[]) {
  std::ifstream in("/Users/tolemy/Downloads/D/in/testcase_06");
  std::cin.rdbuf(in.rdbuf());

  int n, k;
  cin >> n >> k;

  string s;
  cin >> s;

  // ０、もしくは１が連続する区間の長さを管理
  vector<int> ranges0;
  vector<int> ranges1;

  // 後述の計算を簡単にするため、ダミーの区間を追加
  if (s[0] == '0') {
    ranges1.push_back(0);
  }

  {

    int i, j = 0;
    for (i = 0, j = i + 1; j < s.length(); j++) {
      if (s[i] != s[j]) {
        if (s[i] == '0') {

          ranges0.push_back(j - i);
        } else {
          ranges1.push_back(j - i);
        }

        i = j;
      }
    }

    // 最後の区間を追加
    if (s[s.length() - 1] == '0') {

      ranges0.push_back(j - i);
    } else {
      ranges1.push_back(j - i);
    }
  }

  // 後述の計算を簡単にするため、ダミーの区間を追加
  if (s[n - 1] == '0') {
    ranges1.push_back(0);
  }

  // 0の区間をk個含むように連続する区間を計算
  // 0の区間を最大まで含むように合計するには
  // 1の区間1,0の区間1,1の区間2,0の区間2,...,1の区間k,0の区間k,1の区間(k+1)
  // の2k+1個の区間を合計するのが最大。
  // ただし、sが0から始まる場合はの区間1が、0で終わる場合は1の区間(k+1)がない可能性があるため
  // 最初と最後にダミーを追加する

  // 先に最初の2k+1個の区間を計算しておく
  long long int total = 0;
  {
    int i = 0;
    int ubi = min(k, (int)ranges0.size());
    for (i = 0; i < ubi; i++) {
      total += ranges1[i];
      total += ranges0[i];
    }

    total += ranges1[ubi];
  }

  // 0の区間の個数がk個以下の場合は全てを0に変えることができる
  if ((int)ranges0.size() <= k) {
    cout << total << endl;
    return 0;
  }

  long long int total_max = total;
  for (int i = 0; i < ranges0.size() - k; i++) {
    // i=0の時、total = r1[0] + r0[0] + r1[1] + r0[1] + ... + r1[k-1] + r0[k-1]
    // + r1[k] i=1の時、total =                 r1[1] + r0[1] + ... + r1[k-1] +
    // r0[k-1] + r1[k] + r0[k] + r1[k+1]
    // なので、次の和を計算するときは最初の２項を引いて最後の２項を足す
    total -= ranges1[i];
    total -= ranges0[i];
    total += ranges0[i + k];
    total += ranges1[i + k + 1];

    if (total_max < total)
      total_max = total;
  }

  cout << total_max << endl;

  return 0;
}


Fixed submission:


#include <algorithm>
#include <fstream>
#include <functional>
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

#define rep(i, start, end) for (int i = (int)(start); i < (int)(end); i++)
#define charToInt(c) c - 'a'

int main(int argc, const char *argv[]) {

  int n, k;
  cin >> n >> k;

  string s;
  cin >> s;

  // ０、もしくは１が連続する区間の長さを管理
  vector<int> ranges0;
  vector<int> ranges1;

  // 後述の計算を簡単にするため、ダミーの区間を追加
  if (s[0] == '0') {
    ranges1.push_back(0);
  }

  {

    int i, j = 0;
    for (i = 0, j = i + 1; j < s.length(); j++) {
      if (s[i] != s[j]) {
        if (s[i] == '0') {

          ranges0.push_back(j - i);
        } else {
          ranges1.push_back(j - i);
        }

        i = j;
      }
    }

    // 最後の区間を追加
    if (s[s.length() - 1] == '0') {

      ranges0.push_back(j - i);
    } else {
      ranges1.push_back(j - i);
    }
  }

  // 後述の計算を簡単にするため、ダミーの区間を追加
  if (s[n - 1] == '0') {
    ranges1.push_back(0);
  }

  // 0の区間をk個含むように連続する区間を計算
  // 0の区間を最大まで含むように合計するには
  // 1の区間1,0の区間1,1の区間2,0の区間2,...,1の区間k,0の区間k,1の区間(k+1)
  // の2k+1個の区間を合計するのが最大。
  // ただし、sが0から始まる場合はの区間1が、0で終わる場合は1の区間(k+1)がない可能性があるため
  // 最初と最後にダミーを追加する

  // 先に最初の2k+1個の区間を計算しておく
  long long int total = 0;
  {
    int i = 0;
    int ubi = min(k, (int)ranges0.size());
    for (i = 0; i < ubi; i++) {
      total += ranges1[i];
      total += ranges0[i];
    }

    total += ranges1[ubi];
  }

  // 0の区間の個数がk個以下の場合は全てを0に変えることができる
  if ((int)ranges0.size() <= k) {
    cout << total << endl;
    return 0;
  }

  long long int total_max = total;
  for (int i = 0; i < ranges0.size() - k; i++) {
    // i=0の時、total = r1[0] + r0[0] + r1[1] + r0[1] + ... + r1[k-1] + r0[k-1]
    // + r1[k] i=1の時、total =                 r1[1] + r0[1] + ... + r1[k-1] +
    // r0[k-1] + r1[k] + r0[k] + r1[k+1]
    // なので、次の和を計算するときは最初の２項を引いて最後の２項を足す
    total -= ranges1[i];
    total -= ranges0[i];
    total += ranges0[i + k];
    total += ranges1[i + k + 1];

    if (total_max < total)
      total_max = total;
  }

  cout << total_max << endl;

  return 0;
}