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

#include <algorithm>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <functional>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>

using namespace std;

#define LOG(...) printf(__VA_ARGS__)
// #define LOG(...)
#define FOR(i, a, b) for (int i = (int)(a); i < (int)(b); ++i)
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
#define ALL(a) (a).begin(), (a).end()
#define RALL(a) (a).rbegin(), (a).rend()
#define EXIST(s, e) ((s).find(e) != (s).end())
#define SORT(c) sort((c).begin(), (c).end())
#define RSORT(c) sort((c).rbegin(), (c).rend())
#define CLR(a) memset((a), 0, sizeof(a))

typedef long long ll;
typedef unsigned long long ull;
typedef vector<bool> vb;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef vector<vb> vvb;
typedef vector<vi> vvi;
typedef vector<vll> vvll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int dx[] = {-1, 0, 1, 0};
const int dy[] = {0, 1, 0, -1};

struct UnionFind {
  vector<int> v;
  UnionFind(int n) : v(n) {
    for (int i = 0; i < n; i++)
      v[i] = i;
  }
  int find(int x) { return v[x] == x ? x : v[x] = find(v[x]); }
  void unite(int x, int y) { v[find(x)] = find(y); }
};

int caves[1200][1200];
int cx[1100001];
int cy[1100001];
int main() {
  set<int> prime;
  bool b[1000001] = {false};
  FOR(i, 2, 1000001) {
    if (!b[i]) {
      prime.insert(i);
      for (int j = i + i; j < 1000001; j += i)
        b[j] = true;
    }
  }
  int x = 600, y = 600, dir = 2, go = 1, now = 0;
  FOR(i, 1, 1001 * 1001 + 1) {
    caves[y][x] = i;
    cx[i] = x;
    cy[i] = y;
    x += dx[dir];
    y += dy[dir];
    now++;
    if (go == now) {
      if (dir % 2 == 1)
        go++;
      now = 0;
      dir = (dir + 1) % 4;
    }
  }

  int m, n;
  while (cin >> m >> n, m | n) {
    map<int, int> vis;
    queue<int> Q;
    int ans = 0;
    int last = 0;
    if (EXIST(prime, n)) {
      ans = 1;
      last = n;
    }
    Q.push(n);
    while (!Q.empty()) {
      int q = Q.front();
      Q.pop();
      if (EXIST(vis, q))
        continue;
      int sx = cx[q];
      int sy = cy[q];
      int ma = 0;
      REP(i, 3)
      if (EXIST(vis, caves[sy - 1][sx - 1 + i]))
        ma = max(ma, vis[caves[sy - 1][sx - 1 + i]]);
      if (EXIST(prime, q))
        ma++;
      vis[q] = ma;
      if (ans < ma) {
        ans = ma;
        last = q;
      }
      if (ans == ma && EXIST(prime, q)) {
        last = max(last, q);
      }
      REP(i, 3) {
        if (caves[sy + 1][sx - 1 + i] <= m)
          Q.push(caves[sy + 1][sx - 1 + i]);
      }
    }
    cout << ans << ' ' << last << endl;
  }
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <functional>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>

using namespace std;

#define LOG(...) printf(__VA_ARGS__)
// #define LOG(...)
#define FOR(i, a, b) for (int i = (int)(a); i < (int)(b); ++i)
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
#define ALL(a) (a).begin(), (a).end()
#define RALL(a) (a).rbegin(), (a).rend()
#define EXIST(s, e) ((s).find(e) != (s).end())
#define SORT(c) sort((c).begin(), (c).end())
#define RSORT(c) sort((c).rbegin(), (c).rend())
#define CLR(a) memset((a), 0, sizeof(a))

typedef long long ll;
typedef unsigned long long ull;
typedef vector<bool> vb;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef vector<vb> vvb;
typedef vector<vi> vvi;
typedef vector<vll> vvll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int dx[] = {-1, 0, 1, 0};
const int dy[] = {0, 1, 0, -1};

struct UnionFind {
  vector<int> v;
  UnionFind(int n) : v(n) {
    for (int i = 0; i < n; i++)
      v[i] = i;
  }
  int find(int x) { return v[x] == x ? x : v[x] = find(v[x]); }
  void unite(int x, int y) { v[find(x)] = find(y); }
};
set<int> prime;
int left(int a) {
  if (a == 1)
    return 6;
  int loop;
  for (int i = 1; i <= 1000; i += 2)
    if (a <= i * i) {
      loop = i;
      break;
    }
  a -= (loop - 2) * (loop - 2) + 1;
  if (a == 0)
    return a + (loop - 2) * (loop - 2);
  if (0 < a && a < (loop - 2))
    return a + (loop - 2) * (loop - 2) + 1 - (1 + ((loop - 2) / 2) * 8);
  if ((loop - 2) <= a && a < (loop - 2) + loop - 1)
    return a + (loop - 2) * (loop - 2) + 1 + 1;
  if ((loop - 2) + loop - 1 <= a && a < (loop - 2) + loop + loop - 1)
    return a + (loop - 2) * (loop - 2) + 1 + (5 + (loop / 2) * 8);
  if ((loop - 2) + loop + loop - 1 <= a &&
      a <= (loop - 2) + loop + loop - 1 + loop - 1)
    return a + (loop - 2) * (loop - 2) + 1 - 1;
  return -1;
}
int right(int a) {
  if (a == 1)
    return 2;
  int loop;
  for (int i = 1; i <= 1000; i += 2)
    if (a <= i * i) {
      loop = i;
      break;
    }
  a -= (loop - 2) * (loop - 2) + 1;
  if (a < loop - 1)
    return a + (loop - 2) * (loop - 2) + 1 + (1 + (loop / 2) * 8);
  if (a < loop - 1 + loop - 1)
    return a + (loop - 2) * (loop - 2);
  if (a <= loop - 1 + loop - 1 + loop - 3)
    return a + (loop - 2) * (loop - 2) + 1 - (5 + (((loop - 2) / 2) * 8));
  return a + (loop - 2) * (loop - 2) + 1 + 1;
}
int down(int a) {
  if (a == 1)
    return 8;
  int loop;
  for (int i = 1; i <= 1001; i += 2)
    if (a <= i * i) {
      loop = i;
      break;
    }
  a -= (loop - 2) * (loop - 2) + 1;
  if (a == 0)
    return loop * loop;
  if (a < loop - 1)
    return a + (loop - 2) * (loop - 2);
  if (a < loop - 1 + loop - 2)
    return a + (loop - 2) * (loop - 2) + 1 - (3 + (((loop - 2) / 2) * 8));
  if (a < loop - 1 + loop - 2 + loop - 1)
    return a + (loop - 2) * (loop - 2) + 1 + 1;
  return a + (loop - 2) * (loop - 2) + 1 + (7 + (loop / 2) * 8);
}
int main() {
  bool b[1000001] = {false};
  FOR(i, 2, 1000001) {
    if (!b[i]) {
      prime.insert(i);
      for (int j = i + i; j < 1000001; j += i)
        b[j] = true;
    }
  }
  int m, n;
  while (cin >> m >> n, m | n) {
    map<int, int> vis;
    queue<pii> Q;
    int ans = 0;
    int last = 0;
    if (EXIST(prime, n)) {
      ans = 1;
      last = n;
    }
    vis[n] = ans;
    Q.push({n, ans});
    while (!Q.empty()) {
      pii q = Q.front();
      Q.pop();
      if (q.second > ans) {
        ans = q.second;
        last = q.first;
      }
      if (q.second == ans) {
        if (EXIST(prime, q.first))
          last = max(last, q.first);
      }
      REP(i, 3) {
        int q2;
        if (i == 0)
          q2 = down(q.first);
        if (i == 1)
          q2 = left(down(q.first));
        if (i == 2)
          q2 = right(down(q.first));
        if (q2 > m)
          continue;
        if (!EXIST(vis, q2))
          vis[q2] = -1;
        int pr = 0;
        if (EXIST(prime, q2))
          pr = 1;
        if (vis[q2] < q.second + pr) {
          vis[q2] = q.second + pr;
          Q.push({q2, q.second + pr});
        }
      }
    }
    cout << ans << ' ' << last << endl;
  }
}