Problem p02635 (C++)

Buggy submission (status: Runtime Error):
#include <cstdio>
#include <cstring>
const long long mod = 998244353;
char s[400];
long long f[400][400][400];
int count[400], sum[400];
long long preSum[400][400];
const int offset = 310;
long long diaSum[800][400];

int max(int a, int b) { return a > b ? a : b; }

void preWork(int p, int n, int m) {
  for (int i = 0; i <= n; i++) {
    preSum[i][0] = f[p][i][0];
    for (int j = 1; j <= m; j++) {
      preSum[i][j] = preSum[i][j - 1] + f[p][i][j];
      if (preSum[i][j] >= mod)
        preSum[i][j] -= mod;
    }
  }
  for (int i = -m; i <= n; i++) {
    int pos = i + offset;
    if (i >= 0)
      diaSum[pos][0] = f[p][i][0];
    else
      diaSum[pos][0] = 0;
    for (int j = 1; j <= m; j++) {
      diaSum[pos][j] = diaSum[pos][j - 1];
      int last = i + j;
      if (last <= n && last >= 0)
        diaSum[pos][j] += f[p][last][j];
      if (diaSum[pos][j] >= mod)
        diaSum[pos][j] -= mod;
    }
  }
}

long long getPreSum(int i, int j, int n, int m) {
  if (i > n)
    return 0;
  else {
    if (j > m)
      j = m;
    return preSum[i][j];
  }
}

int min(int a, int b) { return a > b ? b : a; }

long long getDiaSum(int i, int j, int n, int m) {
  if (i > n || i < (-m))
    return 0;
  else {
    if (j > m)
      j = m;
    return diaSum[i + offset][j];
  }
}

int main() {
  int m;
  scanf("%s%d", s + 1, &m);
  int len = strlen(s + 1);
  s[++len] = '0';
  int p = 1, size = 0;
  sum[0] = 0;
  while (p <= len) {
    int np = p;
    while (s[np] != '0')
      np++;
    size++;
    count[size] = np - p;
    sum[size] = sum[size - 1] + count[size];
    p = np + 1;
  }
  // for(int i = 1; i <= size; i++) printf("%d ",count[i]);
  // printf("\n");
  f[0][0][0] = 1;
  preWork(0, m, sum[size]);
  for (int i = 1; i <= size; i++) {
    for (int j = 0; j <= m; j++) {
      for (int k = 0; k <= sum[size]; k++) {
        f[i][j][k] = 0;
        if (k >= sum[i]) {
          if (k - count[i] >= 0) {
            long long add =
                getDiaSum(j + count[i] - k, k - count[i], m, sum[size]);
            f[i][j][k] = f[i][j][k] + add;
            if (f[i][j][k] >= mod)
              f[i][j][k] -= mod;
          }
          if (count[i] != 0) {
            long long add = getPreSum(j, k, m, sum[size]) -
                            getPreSum(j, k - count[i], m, sum[size]);
            if (add < 0)
              add += mod;
            f[i][j][k] = f[i][j][k] + add;
            if (f[i][j][k] >= mod)
              f[i][j][k] -= mod;
          }
        }
        // if(f[i][j][k]<0) printf("f[%d][%d][%d] = %lld\n",i,j,k,f[i][j][k]);
      }
    }
    preWork(i, m, sum[size]);
  }
  long long ans = 0;
  for (int i = 0; i <= m; i++)
    ans = (ans + f[size][i][sum[size]]) % mod;
  printf("%lld\n", ans);
  return 0;
}

/*for(int l = 0; l <= k-count[i]; l++){
    int u = k-l;
    int lastJ = j-max(u-count[i],0);
    if(lastJ>=0) f[i][j][k] = (f[i][j][k] + f[i-1][lastJ][l])%mod;
}*/


Fixed submission:
#include <cstdio>
#include <cstring>
const long long mod = 998244353;
char s[400];
long long f[400][400][400];
int count[400], sum[400];
long long preSum[400][400];
const int offset = 310;
long long diaSum[800][400];

int max(int a, int b) { return a > b ? a : b; }

void preWork(int p, int n, int m) {
  for (int i = 0; i <= n; i++) {
    preSum[i][0] = f[p][i][0];
    for (int j = 1; j <= m; j++) {
      preSum[i][j] = preSum[i][j - 1] + f[p][i][j];
      if (preSum[i][j] >= mod)
        preSum[i][j] -= mod;
    }
  }
  for (int i = -m; i <= n; i++) {
    int pos = i + offset;
    if (i >= 0)
      diaSum[pos][0] = f[p][i][0];
    else
      diaSum[pos][0] = 0;
    for (int j = 1; j <= m; j++) {
      diaSum[pos][j] = diaSum[pos][j - 1];
      int last = i + j;
      if (last <= n && last >= 0)
        diaSum[pos][j] += f[p][last][j];
      if (diaSum[pos][j] >= mod)
        diaSum[pos][j] -= mod;
    }
  }
}

long long getPreSum(int i, int j, int n, int m) {
  if (i > n)
    return 0;
  else {
    if (j > m)
      j = m;
    return preSum[i][j];
  }
}

int min(int a, int b) { return a > b ? b : a; }

long long getDiaSum(int i, int j, int n, int m) {
  if (i > n || i < (-m))
    return 0;
  else {
    if (j > m)
      j = m;
    return diaSum[i + offset][j];
  }
}

int main() {
  int m;
  scanf("%s%d", s + 1, &m);
  int len = strlen(s + 1);
  s[++len] = '0';
  int p = 1, size = 0;
  sum[0] = 0;
  while (p <= len) {
    int np = p;
    while (s[np] != '0')
      np++;
    size++;
    count[size] = np - p;
    sum[size] = sum[size - 1] + count[size];
    p = np + 1;
  }
  // for(int i = 1; i <= size; i++) printf("%d ",count[i]);
  // printf("\n");
  m = min(301, m);
  f[0][0][0] = 1;
  preWork(0, m, sum[size]);
  for (int i = 1; i <= size; i++) {
    for (int j = 0; j <= m; j++) {
      for (int k = 0; k <= sum[size]; k++) {
        f[i][j][k] = 0;
        if (k >= sum[i]) {
          if (k - count[i] >= 0) {
            long long add =
                getDiaSum(j + count[i] - k, k - count[i], m, sum[size]);
            f[i][j][k] = f[i][j][k] + add;
            if (f[i][j][k] >= mod)
              f[i][j][k] -= mod;
          }
          if (count[i] != 0) {
            long long add = getPreSum(j, k, m, sum[size]) -
                            getPreSum(j, k - count[i], m, sum[size]);
            if (add < 0)
              add += mod;
            f[i][j][k] = f[i][j][k] + add;
            if (f[i][j][k] >= mod)
              f[i][j][k] -= mod;
          }
        }
        // if(f[i][j][k]<0) printf("f[%d][%d][%d] = %lld\n",i,j,k,f[i][j][k]);
      }
    }
    preWork(i, m, sum[size]);
  }
  long long ans = 0;
  for (int i = 0; i <= m; i++)
    ans = (ans + f[size][i][sum[size]]) % mod;
  printf("%lld\n", ans);
  return 0;
}

/*for(int l = 0; l <= k-count[i]; l++){
    int u = k-l;
    int lastJ = j-max(u-count[i],0);
    if(lastJ>=0) f[i][j][k] = (f[i][j][k] + f[i-1][lastJ][l])%mod;
}*/