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

#include <bits/stdc++.h>
using namespace std;
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
#define FOR(i, c)                                                              \\
  for (__typeof((c).begin()) i = (c).begin(); i != (c).end(); ++i)
#define ALL(c) (c).begin(), (c).end()
#define valid(y, x, h, w) (0 <= y && y < h && 0 <= x && x < w)
#define tpl(...) make_tuple(__VA_ARGS__)
const int INF = 0x3f3f3f3f;
const double EPS = 1e-8;
const double PI = acos(-1);
typedef long long ll;
typedef pair<int, int> pii;
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 (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> ostream &operator<<(ostream &o, const vector<T> &t) {
  o << '[';
  FOR(i, t) {
    if (i != t.begin())
      o << ',';
    o << *i;
  }
  return o << ']';
}
template <class S, class T>
ostream &operator<<(ostream &o, const pair<S, T> &t) {
  return o << '(' << t.first << ',' << t.second << ')';
}
template <int N, class Tp> void out(ostream &, const Tp &) {}
template <int N, class Tp, class, class... Ts>
void out(ostream &o, const Tp &t) {
  if (N)
    o << ',';
  o << get<N>(t);
  out<N + 1, Tp, Ts...>(o, t);
}
template <class... Ts> ostream &operator<<(ostream &o, const tuple<Ts...> &t) {
  o << '(';
  out<0, tuple<Ts...>, Ts...>(o, t);
  return o << ')';
}
template <class T> void output(T *a, int n) {
  REP(i, n) {
    if (i)
      cout << ',';
    cout << a[i];
  }
  cout << endl;
}
template <class T> void output(T *a, int n, int m) {
  REP(i, n) output(a[i], m);
}
template <class T> void output(T t) {
  if (t < 0)
    t = -t, putchar(45);
  static int c[20];
  int k = 0;
  while (t)
    c[k++] = t % 10, t /= 10;
  for (k || (c[k++] = 0); k; putchar(c[--k] ^ 48))
    ;
}
template <class S, class... T> void output(S a, T... t) {
  output(a);
  putchar(' ');
  output(t...);
}
template <class T> bool input(T &t) {
  int n = 1, c;
  for (t = 0; !isdigit(c = getchar()) && ~c && c - 45;)
    ;
  if (!~c)
    return 0;
  for (c - 45 && (n = 0, t = c ^ 48); isdigit(c = getchar());
       t = 10 * t + c - 48)
    ;
  t = n ? -t : t;
  return 1;
}
template <class S, class... T> bool input(S &a, T &...t) {
  input(a);
  return input(t...);
}

ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }

ll modmul(ll a, ll b, ll m) {
  ll r = 0;
  for (; b; b >>= 1, a = (a << 1) % m)
    if (b & 1)
      r = (r + a) % m;
  return r;
}

ll modpow(ll a, ll n, ll m) {
  ll r = 1;
  for (ll x = a % m; n; n >>= 1, x = modmul(x, x, m))
    if (n & 1)
      r = modmul(r, x, m);
  return r;
}

bool Miller_Rabin(ll n, int k = 20) {
  if (n <= 1)
    return 0;
  if (!(n & 1))
    return n == 2;
  if (n % 3 == 0)
    return n == 3;

  ll s = 0, d = n - 1; // n-1 = s^2 * d (d: odd)
  while (!(d & 1))
    s++, d >>= 1;
  while (k--) {
    ll a = (ll)rand() * rand() % (n - 3) + 2;
    ll x = modpow(a, d, n);
    if (x == 1 || x == n - 1)
      continue;
    bool b = 0;
    for (int r = 1; r < s; ++r) {
      x = modmul(x, x, n);
      if (x == 1)
        return 0;
      if (x == n - 1) {
        b = 1;
        break;
      }
    }
    if (!b)
      return 0;
  }
  return 1;
}

int main() {
  int n;
  while (cin >> n) {
    int res = 0;
    REP(i, n) {
      ll A;
      cin >> A;
      A = 2 * A + 1;
      while (A % 2 == 0)
        A /= 2;
      if (A == 1 || Miller_Rabin(A))
        res++;
    }
    cout << res << endl;
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
#define FOR(i, c)                                                              \\
  for (__typeof((c).begin()) i = (c).begin(); i != (c).end(); ++i)
#define ALL(c) (c).begin(), (c).end()
#define valid(y, x, h, w) (0 <= y && y < h && 0 <= x && x < w)
#define tpl(...) make_tuple(__VA_ARGS__)
const int INF = 0x3f3f3f3f;
const double EPS = 1e-8;
const double PI = acos(-1);
typedef long long ll;
typedef pair<int, int> pii;
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 (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> ostream &operator<<(ostream &o, const vector<T> &t) {
  o << '[';
  FOR(i, t) {
    if (i != t.begin())
      o << ',';
    o << *i;
  }
  return o << ']';
}
template <class S, class T>
ostream &operator<<(ostream &o, const pair<S, T> &t) {
  return o << '(' << t.first << ',' << t.second << ')';
}
template <int N, class Tp> void out(ostream &, const Tp &) {}
template <int N, class Tp, class, class... Ts>
void out(ostream &o, const Tp &t) {
  if (N)
    o << ',';
  o << get<N>(t);
  out<N + 1, Tp, Ts...>(o, t);
}
template <class... Ts> ostream &operator<<(ostream &o, const tuple<Ts...> &t) {
  o << '(';
  out<0, tuple<Ts...>, Ts...>(o, t);
  return o << ')';
}
template <class T> void output(T *a, int n) {
  REP(i, n) {
    if (i)
      cout << ',';
    cout << a[i];
  }
  cout << endl;
}
template <class T> void output(T *a, int n, int m) {
  REP(i, n) output(a[i], m);
}
template <class T> void output(T t) {
  if (t < 0)
    t = -t, putchar(45);
  static int c[20];
  int k = 0;
  while (t)
    c[k++] = t % 10, t /= 10;
  for (k || (c[k++] = 0); k; putchar(c[--k] ^ 48))
    ;
}
template <class S, class... T> void output(S a, T... t) {
  output(a);
  putchar(' ');
  output(t...);
}
template <class T> bool input(T &t) {
  int n = 1, c;
  for (t = 0; !isdigit(c = getchar()) && ~c && c - 45;)
    ;
  if (!~c)
    return 0;
  for (c - 45 && (n = 0, t = c ^ 48); isdigit(c = getchar());
       t = 10 * t + c - 48)
    ;
  t = n ? -t : t;
  return 1;
}
template <class S, class... T> bool input(S &a, T &...t) {
  input(a);
  return input(t...);
}

int main() {
  int n;
  while (cin >> n) {
    int res = 0;
    REP(i, n) {
      ll A;
      cin >> A;
      A = 2 * A + 1;
      int cnt = 0;
      for (int k = 2; k * k <= A; ++k) {
        while (A % k == 0) {
          if (k & 1)
            cnt++;
          A /= k;
        }
      }
      if (A > 1)
        cnt++;
      if (cnt < 2)
        res++;
    }
    cout << res << endl;
  }
}