Problem p02635 (C++)

Buggy submission (status: Runtime Error):
#include <iostream>
#include <string>
#include <vector>

using namespace std;

typedef long long ll;

const int MOD = 998244353; // change if needed

struct Modint {
  ll val;

  Modint(ll _val = 0) : val(_val % MOD) {}

  Modint operator+(Modint other) const { return Modint(val + other.val); }

  void operator+=(Modint other) {
    val += other.val;
    val %= MOD;
  }

  Modint operator-() const { return Modint(MOD - val); }

  Modint operator-(Modint other) const { return Modint(val + MOD - other.val); }

  void operator-=(Modint other) {
    val += MOD - other.val;
    val %= MOD;
  }

  Modint operator*(Modint other) const { return Modint(val * other.val); }

  void operator*=(Modint other) {
    val *= other.val;
    val %= MOD;
  }

  bool operator==(Modint other) const { return val == other.val; }

  bool operator!=(Modint other) const { return val != other.val; }
};

Modint exp(Modint a, int k) {
  if (k == 0) {
    return Modint(1);
  } else if (k % 2 == 0) {
    Modint half = exp(a, k / 2);
    return half * half;
  } else {
    return a * exp(a, k - 1);
  }
}

Modint inv(Modint a) { return exp(a, MOD - 2); }

ostream &operator<<(ostream &out, Modint p) {
  out << p.val;
  return out;
}

const int MAX_N = 305;

Modint dp[MAX_N][MAX_N];

Modint prea[MAX_N][MAX_N];
Modint preb[MAX_N][MAX_N];

int main() {
  string s;
  int K;

  cin >> s >> K;

  vector<int> orig;
  int buf = 0;
  for (char c : s) {
    if (c == '0') {
      orig.push_back(buf);
      buf = 0;
    } else {
      buf++;
    }
  }
  orig.push_back(buf);
  buf = 0;

  int n = orig.size();
  int sfx = 0;
  for (int i = 0; i < n; i++) {
    sfx += orig[i];
  }

  dp[0][0] = Modint(1);
  for (int i = 0; i < n; i++) {
    sfx -= orig[i];
    // reset prea and preb
    for (int j = 0; j < MAX_N; j++) {
      for (int k = 0; k < MAX_N; k++) {
        prea[j][k] = 0;
        preb[j][k] = 0;
      }
    }

    // calculate the new prea and preb
    for (int j = 0; j < MAX_N; j++) {
      for (int k = 0; k < MAX_N; k++) {
        prea[j][k] += dp[j][k];
        if (k - orig[i] - 1 >= 0) {
          prea[j][k - orig[i] - 1] -= dp[j][k];
        }

        int lft = sfx - k;
        if (lft > 0) {
          if (j + 1 < MAX_N && k + 1 < MAX_N) {
            preb[j + 1][k + 1] += dp[j][k];
          }

          if (j + lft + 1 < MAX_N && k + lft + 1 < MAX_N) {
            preb[j + lft + 1][k + lft + 1] -= dp[j][k];
          }
        }
      }
    }

    // prefixize prea by k in reverse
    for (int j = 0; j < MAX_N; j++) {
      for (int k = MAX_N - 2; k >= 0; k--) {
        prea[j][k] += prea[j][k + 1];
      }
    }

    // prefixize preb by the diagonals
    for (int j = 1; j < MAX_N; j++) {
      for (int k = 1; k < MAX_N; k++) {
        preb[j][k] += preb[j - 1][k - 1];
      }
    }

    // update the dp
    for (int j = 0; j < MAX_N; j++) {
      for (int k = 0; k < MAX_N; k++) {
        dp[j][k] = prea[j][k] + preb[j][k];
      }
    }
  }

  Modint ans(0);
  for (int j = 0; j <= K; j++) {
    ans += dp[j][0];
  }
  cout << ans << endl;
}


Fixed submission:
#include <iostream>
#include <string>
#include <vector>

using namespace std;

typedef long long ll;

const int MOD = 998244353; // change if needed

struct Modint {
  ll val;

  Modint(ll _val = 0) : val(_val % MOD) {}

  Modint operator+(Modint other) const { return Modint(val + other.val); }

  void operator+=(Modint other) {
    val += other.val;
    val %= MOD;
  }

  Modint operator-() const { return Modint(MOD - val); }

  Modint operator-(Modint other) const { return Modint(val + MOD - other.val); }

  void operator-=(Modint other) {
    val += MOD - other.val;
    val %= MOD;
  }

  Modint operator*(Modint other) const { return Modint(val * other.val); }

  void operator*=(Modint other) {
    val *= other.val;
    val %= MOD;
  }

  bool operator==(Modint other) const { return val == other.val; }

  bool operator!=(Modint other) const { return val != other.val; }
};

Modint exp(Modint a, int k) {
  if (k == 0) {
    return Modint(1);
  } else if (k % 2 == 0) {
    Modint half = exp(a, k / 2);
    return half * half;
  } else {
    return a * exp(a, k - 1);
  }
}

Modint inv(Modint a) { return exp(a, MOD - 2); }

ostream &operator<<(ostream &out, Modint p) {
  out << p.val;
  return out;
}

const int MAX_N = 305;

Modint dp[MAX_N][MAX_N];

Modint prea[MAX_N][MAX_N];
Modint preb[MAX_N][MAX_N];

int main() {
  string s;
  int K;

  cin >> s >> K;

  vector<int> orig;
  int buf = 0;
  for (char c : s) {
    if (c == '0') {
      orig.push_back(buf);
      buf = 0;
    } else {
      buf++;
    }
  }
  orig.push_back(buf);
  buf = 0;

  int n = orig.size();
  int sfx = 0;
  for (int i = 0; i < n; i++) {
    sfx += orig[i];
  }

  dp[0][0] = Modint(1);
  for (int i = 0; i < n; i++) {
    sfx -= orig[i];
    // reset prea and preb
    for (int j = 0; j < MAX_N; j++) {
      for (int k = 0; k < MAX_N; k++) {
        prea[j][k] = 0;
        preb[j][k] = 0;
      }
    }

    // calculate the new prea and preb
    for (int j = 0; j < MAX_N; j++) {
      for (int k = 0; k < MAX_N; k++) {
        prea[j][k] += dp[j][k];
        if (k - orig[i] - 1 >= 0) {
          prea[j][k - orig[i] - 1] -= dp[j][k];
        }

        int lft = sfx - k;
        if (lft > 0) {
          if (j + 1 < MAX_N && k + 1 < MAX_N) {
            preb[j + 1][k + 1] += dp[j][k];
          }

          if (j + lft + 1 < MAX_N && k + lft + 1 < MAX_N) {
            preb[j + lft + 1][k + lft + 1] -= dp[j][k];
          }
        }
      }
    }

    // prefixize prea by k in reverse
    for (int j = 0; j < MAX_N; j++) {
      for (int k = MAX_N - 2; k >= 0; k--) {
        prea[j][k] += prea[j][k + 1];
      }
    }

    // prefixize preb by the diagonals
    for (int j = 1; j < MAX_N; j++) {
      for (int k = 1; k < MAX_N; k++) {
        preb[j][k] += preb[j - 1][k - 1];
      }
    }

    // update the dp
    for (int j = 0; j < MAX_N; j++) {
      for (int k = 0; k < MAX_N; k++) {
        dp[j][k] = prea[j][k] + preb[j][k];
      }
    }
  }

  Modint ans(0);
  for (int j = 0; j <= K && j < MAX_N; j++) {
    ans += dp[j][0];
  }
  cout << ans << endl;
}