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

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;
using namespace placeholders;

using LL = long long;
using ULL = unsigned long long;
using VI = vector<int>;
using VVI = vector<vector<int>>;
using VS = vector<string>;
using ISS = istringstream;
using OSS = ostringstream;
using PII = pair<int, int>;
using VPII = vector<pair<int, int>>;
template <typename T = int> using VT = vector<T>;
template <typename T = int> using VVT = vector<vector<T>>;
template <typename T = int> using LIM = numeric_limits<T>;
template <typename T = int> using OSI = ostream_iterator<T>;

template <typename T> inline istream &operator>>(istream &s, vector<T> &v) {
  for (T &t : v) {
    s >> t;
  }
  return s;
}
template <typename T>
inline ostream &operator<<(ostream &s, const vector<T> &v) {
  for (int i = 0; i < int(v.size()); ++i) {
    s << (" " + !i) << v[i];
  }
  return s;
}

void in_impl(){};
template <typename T, typename... TS> void in_impl(T &head, TS &...tail) {
  cin >> head;
  in_impl(tail...);
}
#define IN(T, ...)                                                             \\
  T __VA_ARGS__;                                                               \\
  in_impl(__VA_ARGS__);

template <typename T> inline T fromString(const string &s) {
  T res;
  istringstream iss(s);
  iss >> res;
  return res;
}
template <typename T> inline string toString(const T &a) {
  ostringstream oss;
  oss << a;
  return oss.str();
}

#define NUMBERED(name, number) NUMBERED2(name, number)
#define NUMBERED2(name, number) name##_##number
#define REP1(n) REP2(NUMBERED(REP_COUNTER, __LINE__), n)
#define REP2(i, n) REP3(i, 0, n)
#define REP3(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
#define GET_REP(a, b, c, F, ...) F
#define REP(...) GET_REP(__VA_ARGS__, REP3, REP2, REP1)(__VA_ARGS__)
#define FOR(e, c) for (auto &&e : c)
#define ALL(c) begin(c), end(c)
#define AALL(a)                                                                \\
  (remove_all_extents<decltype(a)>::type *)a,                                  \\
      (remove_all_extents<decltype(a)>::type *)a +                             \\
          sizeof(a) / sizeof(remove_all_extents<decltype(a)>::type)

#define SZ(v) ((int)(v).size())
#define EXISTS(c, e) ((c).find(e) != (c).end())

template <typename T> inline bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> inline bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

#define PB push_back
#define EM emplace
#define EB emplace_back
#define BI back_inserter

#define MP make_pair
#define fst first
#define snd second

#define DUMP(x) cerr << #x << " = " << (x) << endl

//     Λ   Λ＿_
//  ／(*ﾟーﾟ)／＼
// ／| ￣U U￣|＼／
//   |      |／

int main() {
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
  cout << setprecision(12) << fixed;

  IN(LL, N, X, M);

  VT<LL> counts(M), last(M, -1), done(M);
  LL a = X;
  for (LL i = 0; N;) {
    if (last[a] == -1 || done[a]) {
      ++counts[a];
      last[a] = i;
      (a *= a) %= M;
      --N;
      ++i;
      continue;
    }

    const int d = i - last[a];
    const LL t = N / d;
    LL aa = a;
    do {
      counts[aa] += t;
      (aa *= aa) %= M;
      done[aa] = true;
    } while (aa != a);
    N %= d;
    i += d * t;
  }

  LL res = 0;
  REP(i, M) { res += LL(i) * counts[i]; }
  cout << res << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;
using namespace placeholders;

using LL = long long;
using ULL = unsigned long long;
using VI = vector<int>;
using VVI = vector<vector<int>>;
using VS = vector<string>;
using ISS = istringstream;
using OSS = ostringstream;
using PII = pair<int, int>;
using VPII = vector<pair<int, int>>;
template <typename T = int> using VT = vector<T>;
template <typename T = int> using VVT = vector<vector<T>>;
template <typename T = int> using LIM = numeric_limits<T>;
template <typename T = int> using OSI = ostream_iterator<T>;

template <typename T> inline istream &operator>>(istream &s, vector<T> &v) {
  for (T &t : v) {
    s >> t;
  }
  return s;
}
template <typename T>
inline ostream &operator<<(ostream &s, const vector<T> &v) {
  for (int i = 0; i < int(v.size()); ++i) {
    s << (" " + !i) << v[i];
  }
  return s;
}

void in_impl(){};
template <typename T, typename... TS> void in_impl(T &head, TS &...tail) {
  cin >> head;
  in_impl(tail...);
}
#define IN(T, ...)                                                             \\
  T __VA_ARGS__;                                                               \\
  in_impl(__VA_ARGS__);

template <typename T> inline T fromString(const string &s) {
  T res;
  istringstream iss(s);
  iss >> res;
  return res;
}
template <typename T> inline string toString(const T &a) {
  ostringstream oss;
  oss << a;
  return oss.str();
}

#define NUMBERED(name, number) NUMBERED2(name, number)
#define NUMBERED2(name, number) name##_##number
#define REP1(n) REP2(NUMBERED(REP_COUNTER, __LINE__), n)
#define REP2(i, n) REP3(i, 0, n)
#define REP3(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
#define GET_REP(a, b, c, F, ...) F
#define REP(...) GET_REP(__VA_ARGS__, REP3, REP2, REP1)(__VA_ARGS__)
#define FOR(e, c) for (auto &&e : c)
#define ALL(c) begin(c), end(c)
#define AALL(a)                                                                \\
  (remove_all_extents<decltype(a)>::type *)a,                                  \\
      (remove_all_extents<decltype(a)>::type *)a +                             \\
          sizeof(a) / sizeof(remove_all_extents<decltype(a)>::type)

#define SZ(v) ((int)(v).size())
#define EXISTS(c, e) ((c).find(e) != (c).end())

template <typename T> inline bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> inline bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

#define PB push_back
#define EM emplace
#define EB emplace_back
#define BI back_inserter

#define MP make_pair
#define fst first
#define snd second

#define DUMP(x) cerr << #x << " = " << (x) << endl

//     Λ   Λ＿_
//  ／(*ﾟーﾟ)／＼
// ／| ￣U U￣|＼／
//   |      |／

int main() {
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
  cout << setprecision(12) << fixed;

  IN(LL, N, X, M);

  VT<LL> counts(M), last(M, -1), done(M);
  LL a = X;
  for (int i = 0; a && N; ++i) {
    if (last[a] == -1 || done[a]) {
      ++counts[a];
      last[a] = i;
      (a *= a) %= M;
      --N;
      continue;
    }

    const int d = i - last[a];
    const LL t = N / d;
    LL aa = a;
    do {
      counts[aa] += t;
      (aa *= aa) %= M;
      if (aa == 0) {
        break;
      }
    } while (aa != a);
    N %= d;
    done[a] = true;
  }

  LL res = 0;
  REP(i, M) { res += LL(i) * counts[i]; }
  cout << res << endl;

  return 0;
}