Problem p02635 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 998244353
#define ADD(X, Y) ((X) = ((X) + (Y) % MOD) % MOD)
typedef long long i64;
typedef vector<int> ivec;
typedef vector<string> svec;

char S[330];
int N, K;

i64 dp[2][330][330]; // 0, 1, required K
vector<int> ones;
vector<int> acc;

void clr(int t) {
  for (int i = 0; i <= N; ++i) {
    for (int j = 0; j <= K; ++j)
      dp[t][i][j] = 0;
  }
}

int main() {
  scanf("%s%d", S, &K);

  N = strlen(S);
  S[N++] = '0';
  S[N] = 0;
  int n0 = 0, n1 = 0, n1d = 0;
  for (int i = 0; i < N; ++i) {
    if (S[i] == '0') {
      ++n0;
      ones.push_back(n1d);
      acc.push_back(n1);
      n1d = 0;
    } else {
      ++n1;
      ++n1d;
    }
  }
  K = min(K, N);

  int t = 0;
  clr(t);
  dp[t][0][0] = 1;
  for (int i = 0; i < n0; ++i) {
    for (int j = 0; j <= n1; ++j) {
      for (int k = 0; k <= K; ++k) {
        dp[1 - t][j][k] = dp[t][j][k];
        if (j > 0)
          ADD(dp[1 - t][j][k], dp[1 - t][j - 1][k]);
      }
    }
    for (int j = 0; j <= n1; ++j) {
      for (int k = 0; k <= K; ++k) {
        for (int l = 0; l < ones[i]; ++l) {
          ADD(dp[1 - t][j + l][k], MOD - dp[t][j][k]);
          if (k + ones[i] - l <= K)
            ADD(dp[1 - t][j + l][k + (ones[i] - l)], dp[t][j][k]);
        }
      }
    }
    for (int j = 0; j < acc[i]; ++j) {
      for (int k = 0; k <= K; ++k) {
        dp[1 - t][j][k] = 0;
      }
    }
    t = 1 - t;
  }
  i64 ret = 0;
  for (int i = 0; i <= n1; ++i)
    for (int j = 0; j <= K; ++j)
      ADD(ret, dp[t][i][j]);
  printf("%lld\n", ret);
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 998244353
#define ADD(X, Y) ((X) = ((X) + (Y) % MOD) % MOD)
typedef long long i64;
typedef vector<int> ivec;
typedef vector<string> svec;

char S[330];
int N, K;

i64 dp[2][330][330]; // 0, 1, required K
vector<int> ones;
vector<int> acc;

void clr(int t) {
  for (int i = 0; i <= N; ++i) {
    for (int j = 0; j <= K; ++j)
      dp[t][i][j] = 0;
  }
}

int main() {
  scanf("%s%d", S, &K);

  N = strlen(S);
  S[N++] = '0';
  S[N] = 0;
  int n0 = 0, n1 = 0, n1d = 0;
  for (int i = 0; i < N; ++i) {
    if (S[i] == '0') {
      ++n0;
      ones.push_back(n1d);
      acc.push_back(n1);
      n1d = 0;
    } else {
      ++n1;
      ++n1d;
    }
  }
  K = min(K, N);

  int t = 0;
  clr(t);
  dp[t][0][0] = 1;
  for (int i = 0; i < n0; ++i) {
    for (int j = 0; j <= n1; ++j) {
      for (int k = 0; k <= K; ++k) {
        dp[1 - t][j][k] = dp[t][j][k];
        if (j > 0)
          ADD(dp[1 - t][j][k], dp[1 - t][j - 1][k]);
      }
    }
    for (int j = 0; j <= n1; ++j) {
      for (int k = 0; k <= K; ++k) {
        for (int l = 0; l < ones[i]; ++l) {
          if (j + l > n1)
            break;
          ADD(dp[1 - t][j + l][k], MOD - dp[t][j][k]);
          if (k + ones[i] - l <= K)
            ADD(dp[1 - t][j + l][k + (ones[i] - l)], dp[t][j][k]);
        }
      }
    }
    for (int j = 0; j < acc[i]; ++j) {
      for (int k = 0; k <= K; ++k) {
        dp[1 - t][j][k] = 0;
      }
    }
    t = 1 - t;
  }
  i64 ret = 0;
  for (int i = 0; i <= n1; ++i)
    for (int j = 0; j <= K; ++j)
      ADD(ret, dp[t][i][j]);
  printf("%lld\n", ret);
  return 0;
}