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

//
// 解法PDFを見て解きました（懺悔）
//

#include "bits/stdc++.h"

using namespace std;

using u32 = uint32_t;
using u64 = uint64_t;
using ll = long long;

template <typename U> vector<U> PrimeFactorization(U n) {
#define PrimeFactorization_DebugPrint false

  vector<U> results;
  U a = 1;
#if PrimeFactorization_DebugPrint
  cout << n << " = ";
#endif

  for (U i = 2; i * i <= n; /*NOP*/) {
    auto q = n / i;
    auto r = n - i * q;
    if (r == 0) {
      results.push_back(i);
      n /= i;
#if PrimeFactorization_DebugPrint
      cout << i << " * ";
#endif
    } else {
      i++;
    }
  }

  if (n > 1) {
    results.push_back(n);
#if PrimeFactorization_DebugPrint
    cout << n;
#endif
  }
#if PrimeFactorization_DebugPrint
  cout << '\\n';
#endif

  return results;
}

// コレクションの重複要素を数え、<要素, 個数> のコレクションを返す
template <typename T> map<T, int> countAllElements(const vector<T> &src) {
  map<T, int> results;

  for (auto &e : src) {
    auto it = results.find(e);
    if (it != results.end()) {
      it->second++;
    } else {
      results[e] = 1;
    }
  }

  return results;
}

// map<素数, 個数> からあり得る全ての合成数を列挙する
// - C# の yield があれば配列はいらないのだが。
//   C++ だと boost::coroutines2::coroutine で出来るようだけど今回はなし
template <typename T> vector<T> getCompositeNums(map<T, int> primeNums) {

  vector<T> results;

  if (primeNums.empty()) {
    return results;
  }

  // どの素数を何個ずつ掛け合わせるか
  auto cur(primeNums);
  for (auto &c : cur) {
    c.second = 0;
  }

  auto it = cur.begin();
  while (true) {
    if (++it == cur.end()) {
      u64 com = 1;
      for (auto &c : cur) {
        com *= pow<u64>(c.first, c.second);
#if 0
        cout << c.first << "," << c.second << '\\n';
#endif
      }
      if (com > 1) {
        results.push_back(com);
#if 0
        cout << com << '\\n';
#endif
      }

      bool end = false;
      while (true) {
        if (it == cur.begin()) {
          end = true;
          break;
        }
        it--;
        if (++(it->second) > primeNums[it->first]) {
          it->second = 0;
          continue;
        }
        break;
      }
      if (end) {
        break;
      }
    }
  }

  sort(results.begin(), results.end());
  return results;
}

int solveF(u64 n) {
  auto comps = move(getCompositeNums(countAllElements(PrimeFactorization(n))));
  auto compsM1 =
      move(getCompositeNums(countAllElements(PrimeFactorization(n - 1))));

  // 重複なしコレクション
  set<u64> results;

  // n - 1 の、1以外の自身を含む約数全て
  for (auto k : compsM1) {
    results.insert(k);
  }

  // n の約数それぞれにつき、k で割り切れなくなるまで割り、
  // 割り切れなかったときのあまりが 1 であるか見る
  for (auto k : comps) {
    u64 remainder;
    for (u64 i = n / k, q = i / k; (remainder = i - q * k) == 0;
         i = q, q /= k) {
      /* NOP */
    }
    if (remainder == 1) {
      results.insert(k);
    }
  }

  return results.size();
}

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);

  u64 n;
#if 1
  cin >> n;
  cout << solveF(n);
#elif 0
  // Test Input
  // 2 <= N <= pow(10, 12)
  n = 314159265358;
  cout << solveF(n);
#else
  // Test Input x n
  // for (u64 n = pow(10, 12); n >= 2; n--) {
  for (u64 n = 2; n < pow(10, 12); n++) {
    cout << n << ':' << solveF(n) << '\\n';
  }
#endif

#if 0
  // 約数を列挙するサンプル
  for (auto c : getCompositeNums(countAllElements(PrimeFactorization(12)))) {
    cout << c << ',';
  }
  cout << '\\n';
#endif

  // 以下、解説見る前に色々考えてたやつ（無駄）

  // a^b を pow(a,b)とする

  // n / k^a - bk = 1 ただし a, b は0以上の整数
  // k^a * (bk + 1) = n

  // a = 0:
  // k = (n - 1) / b

  // a = 1:
  // bk^2 + k = n
  // k^2 + k / b = n / b;
  // (k + 0.5 / b)^2 = n / b + 0.25
  // k = sqrt(n / b + 1 / 4) - 0.5 / b
  // k = sqrt(4n / b + 1) / 2 - 1 / 2b
  // k = (sqrt(4n / b + 1) - 1 / b) / 2
  // k が整数なら sqrt(4n / b + 1) - 1 / b は偶数なので、 b = 1が確定. よって
  // k = (sqrt(4n + 1) - 1) / 2

  // a = 2:
  // bk^3 + k^2 = n
  // k^3 + k^2 / b = n / b;
  // (k + 1 / 3b)^3 - (k / 3b^2 + 1 / 27b^3) = n / b
}


Error: TLE

Fixed submission:

//
// 解法PDFを見て解きました（懺悔）
//

#include "bits/stdc++.h"

using namespace std;

using u32 = uint32_t;
using u64 = uint64_t;
using ll = long long;

template <typename U> vector<U> PrimeFactorization(U n) {
#define PrimeFactorization_DebugPrint false

  vector<U> results;
  U a = 1;
#if PrimeFactorization_DebugPrint
  cout << n << " = ";
#endif

  for (U i = 2; i * i <= n; /*NOP*/) {
    auto q = n / i;
    auto r = n - i * q;
    if (r == 0) {
      results.push_back(i);
      n /= i;
#if PrimeFactorization_DebugPrint
      cout << i << " * ";
#endif
    } else {
      i++;
    }
  }

  if (n > 1) {
    results.push_back(n);
#if PrimeFactorization_DebugPrint
    cout << n;
#endif
  }
#if PrimeFactorization_DebugPrint
  cout << '\\n';
#endif

  return results;
}

// コレクションの重複要素を数え、<要素, 個数> のコレクションを返す
template <typename T> map<T, int> countAllElements(const vector<T> &src) {
  map<T, int> results;

  for (auto &e : src) {
    auto it = results.find(e);
    if (it != results.end()) {
      it->second++;
    } else {
      results[e] = 1;
    }
  }

  return results;
}

// map<素数, 個数> からあり得る全ての合成数を列挙する
// - C# の yield があれば配列はいらないのだが。
//   C++ だと boost::coroutines2::coroutine で出来るようだけど今回はなし
template <typename T> vector<T> getCompositeNums(map<T, int> primeNums) {

  vector<T> results;

  // どの素数を何個ずつ掛け合わせるか
  auto cur(primeNums);
  for (auto &c : cur) {
    c.second = 0;
  }

  auto it = cur.begin();
  while (true) {
    if (++it == cur.end()) {
      u64 com = 1;
      for (auto &c : cur) {
        com *= pow(c.first, c.second);
#if 0
        cout << c.first << "," << c.second << '\\n';
#endif
      }
      if (com > 1) {
        results.push_back(com);
#if 0
        cout << com << '\\n';
#endif
      }

      bool end = false;
      while (true) {
        if (it == cur.begin()) {
          end = true;
          break;
        }
        it--;
        if (++(it->second) > primeNums[it->first]) {
          it->second = 0;
          continue;
        }
        break;
      }
      if (end) {
        break;
      }
    }
  }

  sort(results.begin(), results.end());
  return results;
}

int solveF(u64 n) {
  auto comps = move(getCompositeNums(countAllElements(PrimeFactorization(n))));
  auto compsM1 =
      move(getCompositeNums(countAllElements(PrimeFactorization(n - 1))));

  // 重複なしコレクション
  set<u64> results;

  // n - 1 の、1以外の自身を含む約数全て
  for (auto k : compsM1) {
    results.insert(k);
  }

  // n の約数それぞれにつき、k で割り切れなくなるまで割り、
  // 割り切れなかったときのあまりが 1 であるか見る
  for (auto k : comps) {
    u64 remainder;
    for (auto i = n / k, q = i / k; (remainder = i - q * k) == 0;
         i = q, q /= k) {
      /* NOP */
    }
    if (remainder == 1) {
      results.insert(k);
    }
  }

  return results.size();
}

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);

  u64 n;
#if 1
  cin >> n;
  cout << solveF(n);
#elif 0
  // Test Input
  // 2 <= N <= pow(10, 12)
  n = 314159265358;
  cout << solveF(n);
#else
  // Test Input x n
  for (u64 n = pow(10, 12); n >= 2; n--) {
    cout << solveF(n) << '\\n';
  }
#endif

#if 0
  // 約数を列挙するサンプル
  for (auto c : getCompositeNums(countAllElements(PrimeFactorization(12)))) {
    cout << c << ',';
  }
  cout << '\\n';
#endif

  // 以下、解説見る前に色々考えてたやつ（無駄）

  // a^b を pow(a,b)とする

  // n / k^a - bk = 1 ただし a, b は0以上の整数
  // k^a * (bk + 1) = n

  // a = 0:
  // k = (n - 1) / b

  // a = 1:
  // bk^2 + k = n
  // k^2 + k / b = n / b;
  // (k + 0.5 / b)^2 = n / b + 0.25
  // k = sqrt(n / b + 1 / 4) - 0.5 / b
  // k = sqrt(4n / b + 1) / 2 - 1 / 2b
  // k = (sqrt(4n / b + 1) - 1 / b) / 2
  // k が整数なら sqrt(4n / b + 1) - 1 / b は偶数なので、 b = 1が確定. よって
  // k = (sqrt(4n + 1) - 1) / 2

  // a = 2:
  // bk^3 + k^2 = n
  // k^3 + k^2 / b = n / b;
  // (k + 1 / 3b)^3 - (k / 3b^2 + 1 / 27b^3) = n / b
}