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

// http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1189&lang=jp

#include <fstream>
#include <iostream>
#include <math.h>

// #define MAX 1020000

#define MAX 1002

int n, first;
int a[MAX][MAX] = {};
static bool initial;

class Pair {
public:
  int count;
  int last;
  Pair operator+(Pair p) { //	@Override (+)
    this->count += p.count;
    if (p.last != 0)
      this->last = p.last;
    return *this;
  }
  void setPair(int _x, int _y) {
    count = _x;
    last = _y;
  }
};

Pair pair(int a, int b) {
  Pair p;
  p.setPair(a, b);
  return p;
}

bool isPrime(int n) {
  if (n == 1)
    return false;
  if (n == 2)
    return true;
  if (n % 2 == 0)
    return false;

  for (int i = 3; i * i <= n; i += 2) {
    if (n % i == 0)
      return false;
  }
  return true;
}

Pair search(int y, int x) {
  static Pair p[MAX * MAX];

  Pair u_right = pair(-10, 0);
  Pair u_left = pair(-10, 0);

  static bool isReached[MAX * MAX];

  if (initial) { // 初期化
    for (int i = 0; i < MAX * MAX; i++) {
      isReached[i] = false;
      p[i].count = p[i].last = 0;
    }
    initial = false;
  }

  if (n < a[y][x] || a[y][x] == 0)
    return p[0]; // 上限を超えていた場合,,, p[0].count = p[0].last = 0

  // 以下特殊例処理
  if (a[y][x] == 2) {
    p[2].last = 2;
    if (a[y][x] == first)
      initial = true;
    return pair(1, 2) + search(y + 1, x);
  }

  if (n <= 6 && a[y][x] == 3) {
    p[3].last = 2;
    if (a[y][x] == first)
      initial = true;
    return pair(2, 2);
  }
  if (a[y][x] == 4) {
    if (u_right.count == -10)
      u_right = search(y + 1, x + 1);
    if (u_left.count == -10)
      u_left = search(y + 1, x);

    if (u_right.count == u_left.count) {
      if (u_right.last < u_left.last)
        u_right.count--;
    }
    if (u_right.count < u_left.count) {
      isReached[4] = true;
      if (a[y][x] == first)
        initial = true;
      return p[4] = p[4] + u_left;
    } else {
      isReached[4] = true;
      if (a[y][x] == first)
        initial = true;
      return p[4] = p[4] + u_right;
    }
  }
  // ここまで特例処理

  // 偶数の時
  if (a[y][x] % 2 == 0) {
    if (a[y][x] == first) {
      initial = true;
      if (n == first)
        return p[a[y][x]] = search(y + 1, x + 1) +
                            search(y + 1, x - 1); // +p[a[y + 1][x - 1]];
    }
    return p[a[y + 1][x]] = search(y + 1, x);
  }

  // 枝刈り
  if (isReached[a[y][x]]) {
    if (a[y][x] == first)
      initial = true;
    return p[a[y][x]];
  }

  if (u_right.count == -10)
    u_right = search(y + 1, x + 1); // under.rightの計算
  if (u_left.count == -10)
    u_left = search(y + 1, x - 1); // under.leftの計算

  // 素数の時
  if (isPrime(a[y][x])) {
    p[a[y][x]].count = 1;
    p[a[y][x]].last = a[y][x];
  }

  if (u_right.count == u_left.count) {
    if (u_right.last < u_left.last)
      u_right.count--;
  }
  // leftのコストが大きい時
  if (u_right.count < u_left.count) {
    isReached[a[y][x]] = true;
    if (a[y][x] == first)
      initial = true; // 最初に入った時のみ、次のために初期化
    return p[a[y][x]] = p[a[y][x]] + u_left;
  }
  // rightのコストが大きい時
  else {
    isReached[a[y][x]] = true;
    if (a[y][x] == first)
      initial = true; // 最初に入った時のみ、次のために初期化
    return p[a[y][x]] = p[a[y][x]] + u_right;
  }
}

int main() {

  using namespace std;

  int x = MAX / 2;
  int y = MAX / 2;
  int N = 1;
  a[MAX / 2][MAX / 2] = 1;

  for (int i = 0; i < MAX * MAX; i++) {

    for (int n = 0; n < (int)sqrt((double)N); n++) {
      a[y][x] = N;
      N++;
      x++;
    }
    if (N >= (MAX) * (MAX))
      break;
    for (int n = 0; n < (int)sqrt((double)N); n++) {
      a[y][x] = N;
      N++;
      y--;
    }
    if (N >= (MAX) * (MAX))
      break;
    for (int n = 0; n < (int)sqrt((double)N); n++) {
      a[y][x] = N;
      N++;
      x--;
    }
    if (N >= (MAX) * (MAX))
      break;
    for (int n = 0; n < (int)sqrt((double)N); n++) {
      a[y][x] = N;
      N++;
      y++;
    }
    if (N >= (MAX) * (MAX))
      break;
  }

  ifstream fin;
  fin.open("input.txt");

  while (true) {
    cin >> n >> first;
    if (n == 0)
      break;

    initial = true;

    for (int i = 0; i < MAX; i++) {
      for (int j = 0; j < MAX; j++) {
        if (a[i][j] == first) {
          x = j;
          y = i;
          break;
        }
      }
    }

    Pair p = search(y, x);

    cout << p.count << " " << p.last << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

// http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1189&lang=jp

#include <fstream>
#include <iostream>
#include <math.h>

#define MAX 1020000

int n, first;

class Under {
private:
  int right_up;
  int right_down;
  int left_up;
  int left_down;

  int sq;

public:
  int right;
  int straight;
  int left;

  Under(int n) { // constructor

    initialize();

    sq = (int)sqrt(n);
    point(n);

    // 以下の下、右、左、上のじゅんじょ大事。

    if (left_down <= n && n <= right_down) { // 下側にある時
      straight = n + sq / 2 * 8 + 7;
      right = straight + 1;
      left = straight - 1;
    }

    else if (sq * sq + 1 <= n && n <= right_up) { // 右側にある時
      straight = n - 1;
      right = n + 4 * sq + 4;

      if (n == right_down + 1)
        left = n - 2;
      else
        left = n - 4 * sq + 2;
    } else if (sq * sq + 1 <= n && n <= left_down) { // 左側にある時
      straight = n + 1;
      left = n + sq * 4 + 6;

      if (n == sq * sq + sq)
        right = n + 2;
      else
        right = n - 4 * sq + 4;
    } else if (right_up <= n && n <= left_up) { // 上側にある時
      straight = n - (sq + 1) / 2 * 8 + 5;

      right = n - 4 * sq;
      left = n - 4 * sq + 2;
      if (n == sq * sq) {
        left = n + 2;
        right = (sq - 2) * (sq - 2);
      }
      if (n == right_up + 1)
        right = n - 2;
    }
  }

private:
  void point(int n) {

    if (sq % 2) {         // 平方根が奇数
      if (sq * sq == n) { // 平方数
        right_up = n - 3 * sq + 3;
        left_up = n - 2 * sq + 2;
        right_down = n;
        left_down = n - sq + 1;
      } else {
        right_up = n;
        left_up = n;
        right_down = sq * sq + 1;
        left_down = sq * sq - sq + 1;

        if (n < sq * sq + sq + 1)
          left_up = sq * sq - 2 * sq + 2;
        else if (sq * sq + sq + 1 < n)
          right_up = sq * sq + sq + 1;
      }
    }

    else {                // 平方根が偶数
      if (sq * sq == n) { // 平方数
        right_up = n - sq + 1;
        left_up = n;
        right_down = n - 2 * sq + 2;
        left_down = n - 3 * sq + 3;
      } else {
        right_up = sq * sq - sq + 1;
        left_up = sq * sq + 1;
        right_down = n;
        left_down = n;

        if (n < sq * sq + sq + 1)
          right_down = sq * sq - 2 * sq + 2;
        else if (sq * sq + sq + 1 < n)
          left_down = sq * sq + sq + 1;
      }
    }
  }
  void initialize() {
    right = straight = left = right_up = right_down = left_up = left_down = sq =
        0;
  }
};

class Pair {
public:
  int count;
  int last;
  Pair operator+(Pair p) { //	@Override (+)
    this->count += p.count;
    if (p.last != 0)
      this->last = p.last;
    return *this;
  }
  void setPair(int x, int y) {
    count = x;
    last = y;
  }
};

Pair pair(int a, int b) {
  Pair p;
  p.setPair(a, b);
  return p;
}

bool isPrime(int n) {
  if (n == 1)
    return false;
  if (n == 2)
    return true;
  if (n % 2 == 0)
    return false;

  for (int i = 3; i * i <= n; i += 2) {
    if (n % i == 0)
      return false;
  }
  return true;
}

Pair search(int start) {
  static Pair p[MAX];

  Pair right = pair(-10, 0);
  Pair left = pair(-10, 0);

  static bool isReached[MAX];
  static bool initial = true;

  if (initial) { // 初期化
    for (int i = 0; i < MAX; i++) {
      isReached[i] = false;
      p[i].count = p[i].last = 0;
    }
    initial = false;
  }

  if (n < start)
    return p[0]; // 上限を超えていた場合,,, p[0].count = p[0].last = 0

  // 以下特殊例処理
  if (start == 2) {
    p[2].last = 2;
    if (start == first)
      initial = true;
    return pair(1, 2) + search(9);
  }

  if (n <= 6 && start == 3) {
    p[3].last = 2;
    if (start == first)
      initial = true;
    return pair(2, 2);
  }
  if (start == 4) {
    if (right.count == -10)
      right = search(2);
    if (left.count == -10)
      left = search(1);

    if (right.count == left.count) {
      if (right.last < left.last)
        right.count--;
    }
    if (right.count < left.count) {
      //	std::cout << count[4] + left[4] << std::endl;
      isReached[4] = true;
      if (start == first)
        initial = true;
      return p[4] = p[4] + left;
    } else {
      //	std::cout << count[4] + right[4] << std::endl;
      isReached[4] = true;
      if (start == first)
        initial = true;
      return p[4] = p[4] + right;
    }
  }
  // ここまで特例処理

  // 下の数字を計算
  Under under(start);

  // 偶数の時
  if (start % 2 == 0)
    return p[under.straight] = search(under.straight);

  if (right.count == -10)
    right = search(under.right); // under.rightの計算
  if (left.count == -10)
    left = search(under.left); // under.leftの計算

  // std::cout << start << " : l, r = " << left[start] << ", " << right[start]
  // << " : ";

  // 枝刈り
  if (isReached[start]) {
    // std::cout << "\\t"<< count[start] << std::endl;
    return p[start];
  }

  // 素数の時
  if (isPrime(start)) {
    p[start].count = 1;
    p[start].last = start;
  }

  if (right.count == left.count) {
    if (right.last < left.last)
      right.count--;
  }
  // leftのコストが大きい時
  if (right.count < left.count) {
    // std::cout << count[start] + left[start] << std::endl;
    isReached[start] = true;
    if (start == first)
      initial = true; // 最初に入った時のみ、次のために初期化
    return p[start] = p[start] + left;
  }
  // rightのコストが大きい時
  else {
    // std::cout << count[start] + right[start] << std::endl;
    isReached[start] = true;
    if (start == first)
      initial = true; // 最初に入った時のみ、次のために初期化
    return p[start] = p[start] + right;
  }
}

int main() {

  using namespace std;
  /*
          ifstream fin;
          fin.open("input.txt");
  */
  while (true) {
    cin >> n >> first;
    if (n == 0)
      break;

    Pair p = search(first);

    cout << p.count << " " << p.last << endl;
  }
  return 0;
}