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

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <stdlib.h>
#include <time.h>
#define pb push_back
#define IOS                                                                    \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define ll long long
#define ull unsigned long long
#define gcd(a, b) __gcd((a), (b))
#define lcm(a, b) ((a) * (b)) / __gcd((a), (b))
#define clr(x) x.clear()
#define vi vector<int>
#define vll vector<long long>
#define vvi vector<vi>
#define mp make_pair
#define hell 1000000007
#define pii pair<int, int>
#define pll pair<ll, ll>
#define vpii vector<pii>
#define vpll vector<pll>
#define all(a) a.begin(), a.end()
#define ign cin.ignore();

#define o_set                                                                  \\
  tree<int, null_type, less<int>, rb_tree_tag,                                 \\
       tree_order_statistics_node_update>
#define o_setll                                                                \\
  tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update>
// member functions :
// 1. order_of_key(k) : number of elements strictly lesser than k
// 2. find_by_order(k) : k-th element in the set

using namespace std;
using namespace __gnu_pbds;

void __print(int x) { cerr << x; }
void __print(long x) { cerr << x; }
void __print(long long x) { cerr << x; }
void __print(unsigned x) { cerr << x; }
void __print(unsigned long x) { cerr << x; }
void __print(unsigned long long x) { cerr << x; }
void __print(float x) { cerr << x; }
void __print(double x) { cerr << x; }
void __print(long double x) { cerr << x; }
void __print(char x) { cerr << '\\'' << x << '\\''; }
void __print(const char *x) { cerr << '\\"' << x << '\\"'; }
void __print(const string &x) { cerr << '\\"' << x << '\\"'; }
void __print(bool x) { cerr << (x ? "true" : "false"); }

template <typename T, typename V> void __print(const pair<T, V> &x) {
  cerr << '{';
  __print(x.first);
  cerr << ',';
  __print(x.second);
  cerr << '}';
}
template <typename T> void __print(const T &x) {
  int f = 0;
  cerr << '{';
  for (auto &i : x)
    cerr << (f++ ? "," : ""), __print(i);
  cerr << "}";
}
void _print() { cerr << "]\\n"; }
template <typename T, typename... V> void _print(T t, V... v) {
  __print(t);
  if (sizeof...(v))
    cerr << ", ";
  _print(v...);
}
#ifdef MY_PROJECT
#define debug(x...)                                                            \\
  cerr << "[" << #x << "] = [";                                                \\
  _print(x)
#else
#define debug(x...)
#endif

const int N = 1000000;
int lp[N + 1];
vector<int> primes;
void preprocess() {
  for (int i = 2; i <= N; ++i) {
    if (lp[i] == 0) {
      lp[i] = i;
      primes.push_back(i);
    }
    for (int j = 0;
         j < (int)primes.size() && primes[j] <= lp[i] && i * primes[j] <= N;
         ++j)
      lp[i * primes[j]] = primes[j];
  }
}

// ---------- mulmod, powmod and Miller-Rabin test from my toolbox ----------

// return (a*b) % modulo
unsigned long long mulmod(unsigned long long a, unsigned long long b,
                          unsigned long long modulo) {
  // (a * b) % modulo = (a % modulo) * (b % modulo) % modulo
  a %= modulo;
  b %= modulo;

  // fast path
  if (a <= 0xFFFFFFF && b <= 0xFFFFFFF)
    return (a * b) % modulo;

  // we might encounter overflows (slow path)
  // the number of loops depends on b, therefore try to minimize b
  if (b > a)
    std::swap(a, b);

  // bitwise multiplication
  unsigned long long result = 0;
  while (a > 0 && b > 0) {
    // b is odd ? a*b = a + a*(b-1)
    if (b & 1) {
      result += a;
      if (result >= modulo)
        result -= modulo;
      // skip b-- because the bit-shift at the end will remove the lowest bit
      // anyway
    }

    // b is even ? a*b = (2*a)*(b/2)
    a <<= 1;
    if (a >= modulo)
      a -= modulo;

    // next bit
    b >>= 1;
  }

  return result;
}

// return (base^exponent) % modulo => simple implementation
unsigned long long powmod(unsigned long long base, unsigned long long exponent,
                          unsigned long long modulo) {
  unsigned long long result = 1;
  while (exponent > 0) {
    // fast exponentation:
    // odd exponent ? a^b = a*a^(b-1)
    if (exponent & 1)
      result = mulmod(result, base, modulo);

    // even exponent ? a^b = (a*a)^(b/2)
    base = mulmod(base, base, modulo);
    exponent >>= 1;
  }
  return result;
}

// Miller-Rabin-test
bool isPrime(unsigned long long p) {
  // IMPORTANT: requires mulmod(a, b, modulo) and powmod(base, exponent, modulo)

  // some code from
  // https://ronzii.wordpress.com/2012/03/04/miller-rabin-primality-test/ with
  // optimizations from    http://ceur-ws.org/Vol-1326/020-Forisek.pdf good
  // bases can be found at http://miller-rabin.appspot.com/

  // trivial cases
  const unsigned int bitmaskPrimes2to31 =
      (1 << 2) | (1 << 3) | (1 << 5) | (1 << 7) | (1 << 11) | (1 << 13) |
      (1 << 17) | (1 << 19) | (1 << 23) | (1 << 29); // = 0x208A28Ac
  if (p < 31)
    return (bitmaskPrimes2to31 & (1 << p)) != 0;

  if (p % 2 == 0 || p % 3 == 0 || p % 5 == 0 ||
      p % 7 == 0 || // divisible by a small prime
      p % 11 == 0 || p % 13 == 0 || p % 17 == 0)
    return false;

  if (p < 17 * 19) // we filtered all composite numbers < 17*19, all others
                   // below 17*19 must be prime
    return true;

  // fine-tuning for the problem:
  // fast check of more small primes (up to 100)
  if (p % 19 == 0 || p % 23 == 0 || p % 29 == 0 || p % 31 == 0 || p % 37 == 0 ||
      p % 41 == 0 || p % 43 == 0 || p % 47 == 0 || p % 53 == 0 || p % 59 == 0 ||
      p % 61 == 0 || p % 67 == 0 || p % 71 == 0 || p % 73 == 0 || p % 79 == 0 ||
      p % 83 == 0 || p % 89 == 0 || p % 97 == 0)
    return false;

  // test p against those numbers ("witnesses")
  // good bases can be found at http://miller-rabin.appspot.com/
  const unsigned int STOP = 0;
  const unsigned int TestAgainst1[] = {377687, STOP};
  const unsigned int TestAgainst2[] = {31, 73, STOP};
  const unsigned int TestAgainst3[] = {2, 7, 61, STOP};
  // first three sequences are good up to 2^32
  const unsigned int TestAgainst4[] = {2, 13, 23, 1662803, STOP};
  const unsigned int TestAgainst7[] = {2,      325,     9375,       28178,
                                       450775, 9780504, 1795265022, STOP};

  // good up to 2^64
  const unsigned int *testAgainst = TestAgainst7;
  // use less tests if feasible
  if (p < 5329)
    testAgainst = TestAgainst1;
  else if (p < 9080191)
    testAgainst = TestAgainst2;
  else if (p < 4759123141ULL)
    testAgainst = TestAgainst3;
  else if (p < 1122004669633ULL)
    testAgainst = TestAgainst4;

  // find p - 1 = d * 2^j
  auto d = p - 1;
  d >>= 1;
  unsigned int shift = 0;
  while ((d & 1) == 0) {
    shift++;
    d >>= 1;
  }

  // test p against all bases
  do {
    auto x = powmod(*testAgainst++, d, p);
    // is test^d % p == 1 or -1 ?
    if (x == 1 || x == p - 1)
      continue;

    // now either prime or a strong pseudo-prime
    // check test^(d*2^r) for 0 <= r < shift
    bool maybePrime = false;
    for (unsigned int r = 0; r < shift; r++) {
      // x = x^2 % p
      // (initial x was test^d)
      x = mulmod(x, x, p);
      // x % p == 1 => not prime
      if (x == 1)
        return false;

      // x % p == -1 => prime or an even stronger pseudo-prime
      if (x == p - 1) {
        // next iteration
        maybePrime = true;
        break;
      }
    }

    // not prime
    if (!maybePrime)
      return false;
  } while (*testAgainst != STOP);

  // prime
  return true;
}

ll divisors(ll num) {
  ll ans = 1;
  for (ll p : primes) {
    if (p * p * p > num)
      break;
    ll cnt = 1;
    while (num % p == 0) {
      num = num / p;
      cnt++;
    }
    ans = ans * cnt;
  }
  if (isPrime(num)) {
    ans = ans * 2;
    num = 1;
  } else {
    for (ll p : primes) {
      if (num == p * p) {
        ans = ans * 3;
        num = 1;
        break;
      }
    }
  }
  if (num != 1)
    ans = ans * 4;
  return ans;
}

int main() {
  IOS preprocess();
  ll n;
  cin >> n;
  ll res = 1;
  for (ll i = 2; i * i <= n; i++) {
    if (n % i == 0) {
      ll temp = n;
      while (temp % i == 0) {
        temp = temp / i;
      }
      if (temp % i == 1) {
        res++;
      }
    }
  }
  ll num = n - 1;
  ll res2 = divisors(num);
  ll finall = res + res2 - 1;
  cout << finall;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <stdlib.h>
#include <time.h>
#define pb push_back
#define IOS                                                                    \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define ll long long
#define ull unsigned long long
#define gcd(a, b) __gcd((a), (b))
#define lcm(a, b) ((a) * (b)) / __gcd((a), (b))
#define clr(x) x.clear()
#define vi vector<int>
#define vll vector<long long>
#define vvi vector<vi>
#define mp make_pair
#define hell 1000000007
#define pii pair<int, int>
#define pll pair<ll, ll>
#define vpii vector<pii>
#define vpll vector<pll>
#define all(a) a.begin(), a.end()
#define ign cin.ignore();

#define o_set                                                                  \\
  tree<int, null_type, less<int>, rb_tree_tag,                                 \\
       tree_order_statistics_node_update>
#define o_setll                                                                \\
  tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update>
// member functions :
// 1. order_of_key(k) : number of elements strictly lesser than k
// 2. find_by_order(k) : k-th element in the set

using namespace std;
using namespace __gnu_pbds;

void __print(int x) { cerr << x; }
void __print(long x) { cerr << x; }
void __print(long long x) { cerr << x; }
void __print(unsigned x) { cerr << x; }
void __print(unsigned long x) { cerr << x; }
void __print(unsigned long long x) { cerr << x; }
void __print(float x) { cerr << x; }
void __print(double x) { cerr << x; }
void __print(long double x) { cerr << x; }
void __print(char x) { cerr << '\\'' << x << '\\''; }
void __print(const char *x) { cerr << '\\"' << x << '\\"'; }
void __print(const string &x) { cerr << '\\"' << x << '\\"'; }
void __print(bool x) { cerr << (x ? "true" : "false"); }

template <typename T, typename V> void __print(const pair<T, V> &x) {
  cerr << '{';
  __print(x.first);
  cerr << ',';
  __print(x.second);
  cerr << '}';
}
template <typename T> void __print(const T &x) {
  int f = 0;
  cerr << '{';
  for (auto &i : x)
    cerr << (f++ ? "," : ""), __print(i);
  cerr << "}";
}
void _print() { cerr << "]\\n"; }
template <typename T, typename... V> void _print(T t, V... v) {
  __print(t);
  if (sizeof...(v))
    cerr << ", ";
  _print(v...);
}
#ifdef MY_PROJECT
#define debug(x...)                                                            \\
  cerr << "[" << #x << "] = [";                                                \\
  _print(x)
#else
#define debug(x...)
#endif

int main() {
  IOS ll n;
  cin >> n;
  ll res = 1;
  for (ll i = 2; i * i <= n; i++) {
    if (n % i == 0) {
      ll temp = n;
      while (temp % i == 0) {
        temp = temp / i;
      }
      if (temp % i == 1) {
        res++;
      }
    }
  }
  ll num = n - 1;
  ll res2 = 1;
  for (int i = 2; i * i <= num; i++) {
    if (num % i == 0) {
      ll mult = 1;
      while (num % i == 0) {
        num = num / i;
        mult++;
      }
      res2 = res2 * mult;
    }
  }
  if (num > 1)
    res2 = res2 * 2;
  ll finall = res + res2 - 1;
  cout << finall;
  return 0;
}