Problem p02635 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
typedef complex<ld> point;
const int maxn = 300 + 5;
const int mod = 998244353;

int a[maxn];
int dp[maxn][maxn][maxn], dp1[maxn][maxn][maxn], dp2[maxn][maxn][maxn];

int main() {
  ios_base::sync_with_stdio(false);
  string s;
  int K;
  cin >> s >> K;
  int n = 0;
  int cnt = 0;
  for (int i = 0; i < (int)s.size(); i++) {
    if (s[i] == '0') {
      a[n++] = cnt;
      cnt = 0;
    } else
      cnt++;
  }
  a[n++] = cnt;
  int sum = 0;
  for (int i = n; i >= 0; i--) {
    for (int j = 0; j <= 300; j++) {
      for (int k = 0; k <= 300; k++) {
        if (k > sum or k < j)
          dp[i][j][k] = 0;
        else if (i == n)
          dp[i][j][k] = (j == 0 and k == 0);
        else {
          dp[i][j][k] =
              (0ll + dp[i][j][k] + dp1[i + 1][j - 1][j - k + 300] -
               (j - a[i + 1] > 0 ? dp1[i + 1][j - a[i + 1] - 1][j - k + 300]
                                 : 0) +
               mod) %
              mod;
          dp[i][j][k] = (0ll + dp[i][j][k] + dp2[i + 1][300][k] -
                         (j > 0 ? dp2[i + 1][j - 1][k] : 0) + mod) %
                        mod;
          dp[i][j][k] = (0ll + dp[i][j][k] + mod) % mod;
        }

        if (j) {
          dp1[i][j][j - k + 300] =
              (dp1[i][j - 1][j - k + 300] + dp[i][j][k]) % mod;
          dp2[i][j][k] = (dp2[i][j - 1][k] + dp[i][j][k]) % mod;
        } else {
          dp1[i][j][j - k + 300] = dp[i][j][k];
          dp2[i][j][k] = dp[i][j][k];
        }
      }
    }
    sum += a[i];
  }
  int answer = 0;
  for (int j = 0; j <= 300; j++)
    for (int k = 0; k <= min(300, K); k++)
      answer = (answer + dp[0][j][k]) % mod;
  cout << answer << endl;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
typedef complex<ld> point;
const int maxn = 400 + 5;
const int mod = 998244353;

int a[maxn];
ll dp[maxn][maxn][maxn], dp1[maxn][maxn][maxn], dp2[maxn][maxn][maxn];

int main() {
  ios_base::sync_with_stdio(false);
  string s;
  int K;
  cin >> s >> K;
  int n = 0;
  int cnt = 0;
  for (int i = 0; i < (int)s.size(); i++) {
    if (s[i] == '0') {
      a[n++] = cnt;
      cnt = 0;
    } else
      cnt++;
  }
  a[n++] = cnt;
  int sum = 0;
  for (int i = n; i >= 0; i--) {
    for (int j = 0; j <= 400; j++) {
      for (int k = 0; k <= 400; k++) {
        if (k > sum or k < j)
          dp[i][j][k] = 0;
        else if (i == n)
          dp[i][j][k] = (j == 0 and k == 0);
        else {
          dp[i][j][k] =
              (dp[i][j][k] + dp1[i + 1][j - 1][j - k + 400] -
               (j - a[i + 1] > 0 ? dp1[i + 1][j - a[i + 1] - 1][j - k + 400]
                                 : 0) +
               mod) %
              mod;
          dp[i][j][k] = (dp[i][j][k] + dp2[i + 1][400][k] -
                         (j > 0 ? dp2[i + 1][j - 1][k] : 0) + mod) %
                        mod;
          dp[i][j][k] = (dp[i][j][k] + mod) % mod;
        }

        if (j) {
          dp1[i][j][j - k + 400] =
              (dp1[i][j - 1][j - k + 400] + dp[i][j][k]) % mod;
          dp2[i][j][k] = (dp2[i][j - 1][k] + dp[i][j][k]) % mod;
        } else {
          dp1[i][j][j - k + 400] = dp[i][j][k];
          dp2[i][j][k] = dp[i][j][k];
        }
      }
    }
    sum += a[i];
  }
  int answer = 0;
  for (int j = 0; j <= 400; j++)
    for (int k = 0; k <= min(400, K); k++)
      answer = (answer + dp[0][j][k]) % mod;
  cout << answer << endl;
}