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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <vector>
using namespace std;

#define FAST_IO                                                                \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
#define all(x) x.begin(), x.end()
#define ff first
#define ss second
#define pb push_back
#define forn(i, n) for (int i = 0; i < n; ++i)
#define int long long

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;

const int MAXN = 3e5 + 5;
const ll MOD = 1e9 + 7;
const ll INF = INT64_MAX;

random_device rd;
mt19937 rnd(rd());

signed main() {
  FAST_IO;
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++) {
    int cnt = 0;
    for (int x = 1; x * x <= n; x++) {
      for (int y = x; x * x + y * y + x * y <= n; y++) {
        for (int z = y; z * z + x * x + y * y + x * y + y * z + z * x <= n;
             z++) {
          // cout << x << " " << y << " " << z  << " : " << z * z + x * x + y *
          // y + x * y + y * z + z * x << "\\n";
          if (z * z + x * x + y * y + x * y + y * z + z * x == i) {
            set<int> kek;
            kek.insert(x);
            kek.insert(y);
            kek.insert(z);
            if (kek.size() == 1) {
              cnt++;
            }
            if (kek.size() == 2) {
              cnt += 3;
            }
            if (kek.size() == 3) {
              cnt += 6;
            }
          }
        }
      }
    }
    cout << cnt << "\\n";
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <vector>
using namespace std;

#define FAST_IO                                                                \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
#define all(x) x.begin(), x.end()
#define ff first
#define ss second
#define pb push_back
#define forn(i, n) for (int i = 0; i < n; ++i)
#define int long long

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;

const int MAXN = 3e5 + 5;
const ll MOD = 1e9 + 7;
const ll INF = INT64_MAX;

random_device rd;
mt19937 rnd(rd());

signed main() {
  FAST_IO;
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++) {
    int cnt = 0;
    for (int x = 1; x * x <= n; x++) {
      for (int y = 1; x * x + y * y + x * y <= n; y++) {
        for (int z = 1; z * z + x * x + y * y + x * y + y * z + z * x <= n;
             z++) {
          // cout << x << " " << y << " " << z  << " : " << z * z + x * x + y *
          // y + x * y + y * z + z * x << "\\n";
          if (z * z + x * x + y * y + x * y + y * z + z * x == i) {
            cnt++;
          }
        }
      }
    }
    cout << cnt << "\\n";
  }
  return 0;
}