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

#include <bits/stdc++.h>
#define REP(i, n) for (int i = 0; i < (n); i++)
#define FOR(i, s, e) for (int i = s; i < (e); i++)
#define All(v) (v).begin(), (v).end()
#define mp(a, b) make_pair(a, b)
#define pb(a) push_back(a)
#define chmax(x, y) x = max(x, y)
#define chmin(x, y) x = min(x, y)
#define fi first
#define se second
using namespace std;
typedef long long llint;
typedef pair<int, int> P;
const int MOD = (int)1e9 + 7;
const int INF = (int)1e18 * 5;

int a[3000][3000];
int X[1000010], Y[1000010];
int dp[3000][3000];

vector<int> prime;

void PrimeTable(int N) {
  prime = vector<int>(N + 1, 1);
  prime[0] = 0;
  prime[1] = 0;

  for (int i = 2; i * i <= N; ++i) {
    if (prime[i]) {
      for (int j = 2; i * j <= N; ++j) {
        prime[i * j] = 0;
      }
    }
  }
}

void build() {
  int now = 1, dir = 0, diff = 1;
  int x = 1500, y = 1500;
  a[1500][1500] = 1;
  X[1] = 1500, Y[1] = 1500;
  while (now <= 1000000) {
    switch (dir) {
    case 0:
      FOR(dx, 1, diff + 1) {
        a[y][x + dx] = ++now;
        X[now] = x + dx;
        Y[now] = y;
      }
      x = x + diff;
      break;
    case 1:
      FOR(dy, 1, diff + 1) {
        a[y - dy][x] = ++now;
        X[now] = x;
        Y[now] = y - dy;
      }
      y = y - diff;
      break;
    case 2:
      FOR(dx, 1, diff + 1) {
        a[y][x - dx] = ++now;
        X[now] = x - dx;
        Y[now] = y;
      }
      x = x - diff;
      break;
    case 3:
      FOR(dy, 1, diff + 1) {
        a[y + dy][x] = ++now;
        X[now] = x;
        Y[now] = y + dy;
      }
      y = y + diff;
      break;
    }
    if (dir % 2)
      diff++;
    dir = (dir + 1) % 4;
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  build();
  int m, n;
  PrimeTable(1000010);
  cin >> m >> n;
  while (!(m == 0 && n == 0)) {
    REP(i, 3000) REP(j, 3000) dp[i][j] = -INF;
    int sy = Y[n], sx = X[n];
    dp[sy][sx] = prime[n];
    int sf = 1;
    P ans = mp(prime[n], prime[n] ? n : 0);
    bool f = true;
    for (int y = sy; f; y++) {
      f = false;
      FOR(x, sx - sf, sx + sf + 1) {
        if (a[y + 1][x] <= 0 || a[y + 1][x] > m)
          continue;
        f = true;
        dp[y + 1][x] =
            max({dp[y][x - 1], dp[y][x], dp[y][x + 1]}) + prime[a[y + 1][x]];
        if (ans.fi < dp[y + 1][x]) {
          ans.fi = dp[y + 1][x];
          ans.se = a[y + 1][x];
        } else if (ans.fi == dp[y + 1][x]) {
          if (prime[a[y + 1][x]])
            chmax(ans.se, a[y + 1][x]);
        }
      }
      sf++;
    }
    cout << ans.fi << " " << ans.se << "\\n";
    cin >> m >> n;
  }

  return 0;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>
#define REP(i, n) for (int i = 0; i < (n); i++)
#define FOR(i, s, e) for (int i = s; i < (e); i++)
#define All(v) (v).begin(), (v).end()
#define mp(a, b) make_pair(a, b)
#define pb(a) push_back(a)
#define chmax(x, y) x = max(x, y)
#define chmin(x, y) x = min(x, y)
#define fi first
#define se second
#define int long long
using namespace std;
typedef long long llint;
typedef pair<int, int> P;
const int MOD = (int)1e9 + 7;
const int INF = (int)1e18 * 5;

int a[3000][3000];
int X[1000010], Y[1000010];
P dp[3000][3000];

vector<int> prime;

void PrimeTable(int N) {
  prime = vector<int>(N + 1, 1);
  prime[0] = 0;
  prime[1] = 0;

  for (int i = 2; i * i <= N; ++i) {
    if (prime[i]) {
      for (int j = 2; i * j <= N; ++j) {
        prime[i * j] = 0;
      }
    }
  }
}

signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int now = 1, dir = 0, diff = 1;
  int x = 1500, y = 1500;
  a[1500][1500] = 1;
  X[1] = 1500, Y[1] = 1500;
  while (now <= 1000000) {
    switch (dir) {
    case 0:
      FOR(dx, 1, diff + 1) {
        a[y][x + dx] = ++now;
        X[now] = x + dx;
        Y[now] = y;
      }
      x = x + diff;
      break;
    case 1:
      FOR(dy, 1, diff + 1) {
        a[y - dy][x] = ++now;
        X[now] = x;
        Y[now] = y - dy;
      }
      y = y - diff;
      break;
    case 2:
      FOR(dx, 1, diff + 1) {
        a[y][x - dx] = ++now;
        X[now] = x - dx;
        Y[now] = y;
      }
      x = x - diff;
      break;
    case 3:
      FOR(dy, 1, diff + 1) {
        a[y + dy][x] = ++now;
        X[now] = x;
        Y[now] = y + dy;
      }
      y = y + diff;
      break;
    }
    if (dir % 2)
      diff++;
    dir = (dir + 1) % 4;
  }

  // REP(i, 20){
  //     REP(j, 20){
  //         printf("%03d ", a[1490+i][1490+j]);
  //     }
  //     printf("\\n");
  // }
  int m, n;
  cin >> m >> n;
  PrimeTable(1000010);
  while (!(m == 0 && n == 0)) {
    REP(i, 3000) REP(j, 3000) dp[i][j] = mp(-1, 0);
    int ny = Y[n], nx = X[n];
    dp[ny][nx] = mp(prime[n], prime[n] ? n : 0);
    bool f = true;
    vector<int> shiftX{0};
    vector<int> downX{-1, 0, 1};
    int down = 0;
    while (f) {
      f = false;
      for (auto sx : shiftX) {
        for (auto dx : downX) {
          if (dp[ny + 1][nx + sx + dx].fi <= dp[ny][nx + sx].fi) {
            dp[ny + 1][nx + sx + dx].fi =
                dp[ny][nx + sx].fi + prime[a[ny + 1][nx + sx + dx]];
            if (prime[a[ny + 1][nx + sx + dx]])
              dp[ny + 1][nx + sx + dx].se =
                  max(a[ny + 1][nx + sx + dx], dp[ny][nx + sx].se);
            else {
              if (dp[ny + 1][nx + sx + dx].fi < dp[ny][nx + sx].fi)
                dp[ny + 1][nx + sx + dx].se = dp[ny][nx + sx].se;
              else
                chmax(dp[ny + 1][nx + sx + dx].se, dp[ny][nx + sx].se);
            }
          } else if (dp[ny + 1][nx + sx + dx].fi ==
                     dp[ny][nx + sx].fi + prime[a[ny + 1][nx + sx + dx]]) {
            dp[ny + 1][nx + sx + dx].se = a[ny + 1][nx + sx + dx];
          }
          if (0 < a[ny + 1][nx + sx + dx] && a[ny + 1][nx + sx + dx] <= m)
            f = true;
          else
            dp[ny + 1][nx + sx + dx].fi = -1;
        }
      }

      if (down < 1000) {
        down++;
        shiftX.pb(down);
        shiftX.pb(-down);
      }
      // sort(All(shiftX));
      // for(auto sx : shiftX){
      //     cout << dp[ny+1][nx+sx]  << " ";
      // }
      // cout << "\\n";
      ny++;
    }
    P ans = mp(0, 0);
    REP(i, 3000) chmax(ans, dp[ny - 2][i]);
    REP(i, 3000) chmax(ans, dp[ny - 1][i]);
    cout << ans.fi << " " << ans.se << "\\n";

    cin >> m >> n;
  }

  return 0;
}