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

#ifdef _MSC_VER
#include "stdc++.h"
#endif
#ifdef __GNUC__
#include <bits/stdc++.h>
#endif
using namespace std;

using ll = long long;

// vector
using vi = vector<int>;
using vii = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;

// constant
const double EPS = 1e-10;

#define REP(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, s, e) for (ll i = s; i < e; i++)
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()

// debug
#define DUMP(x) cerr << #x << " = " << (x) << endl;
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

// vector
// template<typename T> vector<vector<T>> vv(ll H, ll W, T v) { return
// vector<vector<T>> vec(H, vector<T>(W, v)); }

// view vector
template <typename T> void view(T e) { cout << e << endl; }
template <typename T> void view(const vector<T> &v) {
  for (const auto &e : v) {
    cout << e << " ";
  }
  cout << endl;
}
template <typename T> void view(const vector<vector<T>> &vv) {
  for (const auto &v : vv) {
    view(v);
  }
}

// DP
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}

// ***************************************

ll pw(ll x, ll y, ll z) { return pow(x, 2) + pow(y, 2) + pow(z, 2); }

ll f(ll x, ll y, ll z) { return pw(x, y, z) + x * y + y * z + z * x; }

int main() {

  ll N;
  cin >> N;

  FOR(n, 1, N + 1) {
    ll x, y, z;
    ll cnt = 0;
    for (x = 1; f(x, x, x) <= n; x++) {
      for (y = x; f(x, y, y) <= n; y++) {
        for (z = y; f(x, y, z) <= n; z++) {
          if (pw(x, y, z) >= n)
            break;
          if (f(x, y, z) == n) {
            set<ll> s{x, y, z};
            if (s.size() == 3)
              cnt += 6;
            else if (s.size() == 2)
              cnt += 3;
            else
              cnt++;
          }
        }
        if (z == 1)
          break;
      }
      if (y == 1 && z == 1)
        break;
    }

    cout << cnt << endl;
  }

  return 0;
}


Error: TLE

Fixed submission:

#ifdef _MSC_VER
#include "stdc++.h"
#endif
#ifdef __GNUC__
#include <bits/stdc++.h>
#endif
using namespace std;

using ll = long long;

// vector
using vi = vector<int>;
using vii = vector<vi>;
using vll = vector<ll>;
using vvll = vector<vll>;

// constant
const double EPS = 1e-10;

#define REP(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, s, e) for (ll i = s; i < e; i++)
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()

// debug
#define DUMP(x) cerr << #x << " = " << (x) << endl;
#define DEBUG(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl;

// vector
// template<typename T> vector<vector<T>> vv(ll H, ll W, T v) { return
// vector<vector<T>> vec(H, vector<T>(W, v)); }

// view vector
template <typename T> void view(T e) { cout << e << endl; }
template <typename T> void view(const vector<T> &v) {
  for (const auto &e : v) {
    cout << e << " ";
  }
  cout << endl;
}
template <typename T> void view(const vector<vector<T>> &vv) {
  for (const auto &v : vv) {
    view(v);
  }
}

// DP
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}

// ***************************************

ll pw(ll x, ll y, ll z) { return pow(x, 2) + pow(y, 2) + pow(z, 2); }

ll f(ll x, ll y, ll z) { return pw(x, y, z) + x * y + y * z + z * x; }

int main() {

  ll N;
  cin >> N;

  FOR(n, 1, N + 1) {
    ll x, y, z;
    ll cnt = 0;
    for (x = 1; f(x, 1, 1) <= n; x++) {
      for (y = 1; f(x, y, 1) <= n; y++) {
        ll zend = sqrt(n) + 1;
        if (f(x, y, zend) < n)
          continue;
        for (z = 1; f(x, y, z) <= n; z++) {
          if (pw(x, y, z) >= n)
            break;
          if (f(x, y, z) == n)
            cnt++;
        }
        if (z == 1)
          break;
      }
      if (y == 1 && z == 1)
        break;
    }

    cout << cnt << endl;
  }

  return 0;
}