Problem p02635 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>
#define rep(i, n) for (i = 0; i < n; i++)
using namespace std;

void reverse(string &s) {
  int l = 0, r = s.length() - 1;
  while (l < r) {
    swap(s[l], s[r]);
    l++;
    r--;
  }
}

int mod = 998244353;
string s;
int K;
vector<int> a;
vector<int> ra;
int dp[310][310][310]; // dp[index][bsum][cost] = 何通り？

int main() {
  int i;

  cin >> s >> K;
  reverse(s);

  int len = 0;
  rep(i, s.length()) {
    if (s[i] == '0') {
      a.push_back(len);
      len = 0;
    } else {
      len++;
    }
  }
  a.push_back(len);

  ra.resize(a.size() + 1);
  ra[0] = 0;
  rep(i, a.size()) { ra[i + 1] = ra[i] + a[i]; }

  dp[0][0][0] = 1;
  rep(i, a.size()) {
    for (int bsum = 0; bsum <= ra[i]; bsum++) {
      for (int cst = 0; cst <= ra[i]; cst++) {
        if (dp[i][bsum][cst] == 0)
          continue;
        for (int b = 0; bsum + b <= ra[i + 1]; b++) {
          int plus = a[i] - b;
          if (plus < 0)
            plus = 0;
          dp[i + 1][bsum + b][cst + plus] += dp[i][bsum][cst];
          if (dp[i + 1][bsum + b][cst + plus] >= mod) {
            dp[i + 1][bsum + b][cst + plus] -= mod;
          }
        }
      }
    }
  }

  // dp[a.size()][*][≦K]
  int ans = 0;
  for (int bsum = ra[a.size()]; bsum <= ra[a.size()]; bsum++) {
    for (int cst = 0; cst <= K; cst++) {
      // cout << "dp[" << a.size() << ", " << bsum << ", " << cst << "] = " <<
      // dp[a.size()][bsum][cst] << endl;
      ans += dp[a.size()][bsum][cst];
      ans %= mod;
    }
  }
  cout << ans << endl;
  return 0;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>
#define rep(i, n) for (i = 0; i < n; i++)
using namespace std;

void reverse(string &s) {
  int l = 0, r = s.length() - 1;
  while (l < r) {
    swap(s[l], s[r]);
    l++;
    r--;
  }
}

int mod = 998244353;
string s;
int K;
vector<int> a;
vector<int> ra;
int dp[310][310][310]; // dp[index][bsum][cost] = 何通り？

int main() {
  int i;

  cin >> s >> K;
  reverse(s);
  K = min(K, 301);

  int len = 0;
  rep(i, s.length()) {
    if (s[i] == '0') {
      a.push_back(len);
      len = 0;
    } else {
      len++;
    }
  }
  a.push_back(len);

  ra.resize(a.size() + 1);
  ra[0] = 0;
  rep(i, a.size()) { ra[i + 1] = ra[i] + a[i]; }

  dp[0][0][0] = 1;
  rep(i, a.size()) {
    for (int bsum = 0; bsum <= ra[i]; bsum++) {
      for (int cst = 0; cst <= ra[i]; cst++) {
        if (dp[i][bsum][cst] == 0)
          continue;
        for (int b = 0; bsum + b <= ra[i + 1]; b++) {
          int plus = a[i] - b;
          if (plus < 0)
            plus = 0;
          dp[i + 1][bsum + b][cst + plus] += dp[i][bsum][cst];
          if (dp[i + 1][bsum + b][cst + plus] >= mod) {
            dp[i + 1][bsum + b][cst + plus] -= mod;
          }
        }
      }
    }
  }

  // dp[a.size()][*][≦K]
  int ans = 0;
  for (int bsum = ra[a.size()]; bsum <= ra[a.size()]; bsum++) {
    for (int cst = 0; cst <= K; cst++) {
      // cout << "dp[" << a.size() << ", " << bsum << ", " << cst << "] = " <<
      // dp[a.size()][bsum][cst] << endl;
      ans += dp[a.size()][bsum][cst];
      ans %= mod;
    }
  }
  cout << ans << endl;
  return 0;
}