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

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define rep(i, n) for (int i = 0; i < (n); i++)
#define chmin(a, b)                                                            \\
  if ((a) > (b))                                                               \\
    (a) = (b);
#define chmax(a, b)                                                            \\
  if ((a) < (b))                                                               \\
    (a) = (b);
#define vi vector<int>
#define pii pair<int, int>
#define all(v) (v).begin(), (v).end()
#define allr(v) (v).rbegin(), (v).rend()
#define pb push_back
#define pf push_front
int gcd(int a, int b) { /*a>=0,b>=0,￢(a=b=0)*/
  while (min(a, b) > 0) {
    if (a < b)
      swap(a, b);
    a %= b;
  }
  return max(a, b);
}
int dx[] = {1, 0, -1, 0, 1, -1, -1, 1};
int dy[] = {0, 1, 0, -1, 1, 1, -1, -1};
const int MOD = 1e9 + 7;
const long long INF = 1e18 + 10;

/*--------------------------------------------------------------------*/

signed main() {
  int n;
  cin >> n;

  FOR(i, 1, n) {
    int ans = 0;

    FOR(x, 1, sqrt(i)) {
      FOR(y, x, sqrt(i)) {
        if (x * x + y * y + x * y > n) {
          break;
        }
        int z = 0;
        int ng = 0, ok = n + 1;
        int mid;
        while (ok - ng > 1) {
          mid = (ok + ng) / 2;
          if (2 * i <= (x + y) * (x + y) + (y + mid) * (y + mid) +
                           (mid + x) * (mid + x)) {
            ok = mid;
          } else {
            ng = mid;
          }
        }
        z = ok;
        if (2 * i !=
            (x + y) * (x + y) + (y + z) * (y + z) + (z + x) * (z + x)) {
          continue;
        }

        ans++;
        if (x < y) {
          ans++;
        }
        // printf("%d %d %d\\n",x,y,z);
      }
    }

    cout << ans << endl;
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define rep(i, n) for (int i = 0; i < (n); i++)
#define chmin(a, b)                                                            \\
  if ((a) > (b))                                                               \\
    (a) = (b);
#define chmax(a, b)                                                            \\
  if ((a) < (b))                                                               \\
    (a) = (b);
#define vi vector<int>
#define pii pair<int, int>
#define all(v) (v).begin(), (v).end()
#define allr(v) (v).rbegin(), (v).rend()
#define pb push_back
#define pf push_front
int gcd(int a, int b) { /*a>=0,b>=0,￢(a=b=0)*/
  while (min(a, b) > 0) {
    if (a < b)
      swap(a, b);
    a %= b;
  }
  return max(a, b);
}
int dx[] = {1, 0, -1, 0, 1, -1, -1, 1};
int dy[] = {0, 1, 0, -1, 1, 1, -1, -1};
const int MOD = 1e9 + 7;
const long long INF = 1e18 + 10;

/*--------------------------------------------------------------------*/

signed main() {
  int n;
  cin >> n;

  FOR(i, 1, n) {
    int ans = 0;

    FOR(x, 1, sqrt(n) + 2) {
      FOR(y, x, sqrt(n) + 2) {
        int z = 0;
        int ng = 0, ok = n + 1;
        int mid;
        while (ok - ng > 1) {
          mid = (ok + ng) / 2;
          if (2 * i <= (x + y) * (x + y) + (y + mid) * (y + mid) +
                           (mid + x) * (mid + x)) {
            ok = mid;
          } else {
            ng = mid;
          }
        }
        z = ok;
        if (2 * i !=
            (x + y) * (x + y) + (y + z) * (y + z) + (z + x) * (z + x)) {
          continue;
        }

        ans++;
        if (x < y) {
          ans++;
        }
        // printf("%d %d %d\\n",x,y,z);
      }
    }

    cout << ans << endl;
  }

  return 0;
}