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

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for (ll i = 0; i < (n); ++i)
using ll = int_fast64_t;
using pll = pair<ll, ll>;
constexpr ll INF = (1LL << 60);
constexpr ll MOD = (1e9 + 7);
// constexpr ll MOD = (998244353);
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;
}
#if loc || local
void dump() { cerr << endl; }
template <class T, class... Ts> void dump(T &&h, Ts &&...t) {
  cerr << h << ", ";
  dump(forward<Ts>(t)...);
}
#else
void dump() {}
template <class T, class... Ts> void dump(T &&h, Ts &&...t) {}
#endif
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &elemnt : v)
    is >> elemnt;
  return is;
}
template <class T, class U> istream &operator>>(istream &is, pair<T, U> &p) {
  is >> p.first >> p.second;
  return is;
}
template <class T> ostream &operator<<(ostream &os, vector<T> const &v) {
  for (auto const &vi : v)
    os << vi << " ";
  return os;
}
template <class T, class U>
ostream &operator<<(ostream &os, pair<T, U> const &p) {
  os << p.first << "," << p.second;
  return os;
}
template <class T> vector<T> vec(size_t a) { return vector<T>(a); }
template <class T, class... Ts> auto vec(size_t a, Ts... ts) {
  return vector<decltype(vec<T>(ts...))>(a, vec<T>(ts...));
}

std::vector<ll> divisor(ll n) {
  std::vector<ll> div;
  for (ll p = 1; p * p <= n; ++p) {
    if (n % p == 0) {
      div.push_back(p);
      if (p * p != n)
        div.push_back(n / p);
    }
  }
  std::sort(div.begin(), div.end());
  return div;
}
map<ll, ll> p;

bool is_prime(ll x) {
  if (x == 1)
    return false;
  if (x == 2)
    return false;
  for (ll p = 3; p * p <= x; p += 2) {
    if (x % p == 0)
      return false;
  }
  return true;
}

signed main() {

  ll n;
  cin >> n;

  int ans = 0;
  auto ok = [](ll n, ll k) {
    if (k == 1)
      return false;
    if (n == k)
      return true;
    while (n >= k) {
      if (n % k == 0)
        n /= k;
      else
        n %= k;
    }
    return n == 1;
  };

  for (auto x : divisor(n)) {
    ans += ok(n, x);
  }
  for (auto x : divisor(n - 1)) {
    ans += ok(n, x);
  }
  cout << (ans) << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for (ll i = 0; i < (n); ++i)
using ll = int_fast64_t;
using pll = pair<ll, ll>;
constexpr ll INF = (1LL << 60);
constexpr ll MOD = (1e9 + 7);
// constexpr ll MOD = (998244353);
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;
}
#if loc || local
void dump() { cerr << endl; }
template <class T, class... Ts> void dump(T &&h, Ts &&...t) {
  cerr << h << ", ";
  dump(forward<Ts>(t)...);
}
#else
void dump() {}
template <class T, class... Ts> void dump(T &&h, Ts &&...t) {}
#endif
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &elemnt : v)
    is >> elemnt;
  return is;
}
template <class T, class U> istream &operator>>(istream &is, pair<T, U> &p) {
  is >> p.first >> p.second;
  return is;
}
template <class T> ostream &operator<<(ostream &os, vector<T> const &v) {
  for (auto const &vi : v)
    os << vi << " ";
  return os;
}
template <class T, class U>
ostream &operator<<(ostream &os, pair<T, U> const &p) {
  os << p.first << "," << p.second;
  return os;
}
template <class T> vector<T> vec(size_t a) { return vector<T>(a); }
template <class T, class... Ts> auto vec(size_t a, Ts... ts) {
  return vector<decltype(vec<T>(ts...))>(a, vec<T>(ts...));
}

std::vector<ll> divisor(ll n) {
  std::vector<ll> div;
  for (ll p = 1; p * p <= n; ++p) {
    if (n % p == 0) {
      div.push_back(p);
      if (p * p != n)
        div.push_back(n / p);
    }
  }
  std::sort(div.begin(), div.end());
  return div;
}
map<ll, ll> p;

bool is_prime(ll x) {
  if (x == 1)
    return false;
  if (x == 2)
    return false;
  for (ll p = 3; p * p <= x; p += 2) {
    if (x % p == 0)
      return false;
  }
  return true;
}

void dfs(ll x) {
  if (p.count(x))
    return;
  for (auto pi : divisor(x)) {
    p[pi]++;
    if (!is_prime(pi))
      continue;
    dfs(pi);
  }
}

signed main() {

  ll n;
  cin >> n;

  dfs(n - 1);
  dfs(n);
  int ans = 0;
  auto ok = [](ll n, ll k) {
    if (k == 1)
      return false;
    while (n >= k) {
      if (n % k == 0)
        n /= k;
      else
        n -= k;
    }
    return n == 1;
  };
  for (auto x : p) {
    ans += ok(n, x.first);
  }
  cout << (ans) << endl;
}