nt/count_by_factor_type.hpp

• View this file on GitHub
• Last update: 2024-09-30 22:44:02+09:00
• Include: #include "nt/count_by_factor_type.hpp"
• Link: https://atcoder.jp/contests/xmascon20/tasks/xmascon20_d

Depends on

• nt/primesum.hpp
• nt/primetable.hpp

Verified with

• test/1_mytest/count_by_factor_type.test.cpp

Code

#include "nt/primetable.hpp"
#include "nt/primesum.hpp"

// factor type: 降順 270 -> (3,1,1)
// N=10^9: 1324 種類, 0.4sec
// https://atcoder.jp/contests/xmascon20/tasks/xmascon20_d
map<vc<int>, ll> count_by_factor_type(ll N) {
  ll sqN = sqrtl(N);
  auto P = primetable<int>(sqN);
  PrimeSum<ll> X(N);
  X.calc_count();

  // 1 and prime
  map<vc<int>, ll> ANS;
  ANS[vc<int>()] = 1;
  if (X[N] > 0) ANS[vc<int>({1})] = X[N];

  auto add = [&](vc<int> F, int k) -> vc<int> {
    int p = len(F);
    F.eb(k);
    while (0 < p && F[p - 1] < F[p]) { swap(F[p - 1], F[p]), --p; }
    return F;
  };

  // t = up_i^k のときに
  auto dfs = [&](auto& dfs, ll t, ll i, ll k, vc<int> U) -> void {
    // U * primes を追加する
    vc<int> nxt1 = add(U, k + 1);
    ANS[nxt1]++;
    vc<int> Uk = add(U, k);
    vc<int> nxt2 = add(Uk, 1);
    ll cnt = X[N / t] - X[P[i]];
    if (cnt > 0) ANS[nxt2] += X[N / t] - X[P[i]];
    ll lim = sqrtl(double(N) / t);
    if (P[i] <= lim) { dfs(dfs, t * P[i], i, k + 1, U); }
    FOR(j, i + 1, len(P)) {
      if (P[j] > lim) break;
      dfs(dfs, t * P[j], j, 1, Uk);
    }
  };
  FOR(i, len(P)) if (P[i] <= sqN) dfs(dfs, P[i], i, 1, {});
  return ANS;
}

#line 2 "nt/primetable.hpp"

template <typename T = int>
vc<T> primetable(int LIM) {
  ++LIM;
  const int S = 32768;
  static int done = 2;
  static vc<T> primes = {2}, sieve(S + 1);

  if (done < LIM) {
    done = LIM;

    primes = {2}, sieve.assign(S + 1, 0);
    const int R = LIM / 2;
    primes.reserve(int(LIM / log(LIM) * 1.1));
    vc<pair<int, int>> cp;
    for (int i = 3; i <= S; i += 2) {
      if (!sieve[i]) {
        cp.eb(i, i * i / 2);
        for (int j = i * i; j <= S; j += 2 * i) sieve[j] = 1;
      }
    }
    for (int L = 1; L <= R; L += S) {
      array<bool, S> block{};
      for (auto& [p, idx]: cp)
        for (int i = idx; i < S + L; idx = (i += p)) block[i - L] = 1;
      FOR(i, min(S, R - L)) if (!block[i]) primes.eb((L + i) * 2 + 1);
    }
  }
  int k = LB(primes, LIM + 1);
  return {primes.begin(), primes.begin() + k};
}
#line 3 "nt/primesum.hpp"

/*
N と完全乗法的関数 f の prefix sum 関数 F を与える。
n = floor(N/d) となる n に対する sum_{p <= n} f(p) を計算する。
特に、素数の k 乗和や、mod m ごとでの素数の k 乗和が計算できる。
Complexity: O(N^{3/4}/logN) time, O(N^{1/2}) space.
*/
template <typename T>
struct PrimeSum {
  ll N;
  ll sqN;
  vc<T> sum_lo, sum_hi;
  bool calculated;

  PrimeSum(ll N) : N(N), sqN(sqrtl(N)), calculated(0) {}

  // [1, x] ただし、x = floor(N, i) の形

  T operator[](ll x) {
    assert(calculated);
    return (x <= sqN ? sum_lo[x] : sum_hi[double(N) / x]);
  }

  template <typename F>
  void calc(const F f) {
    auto primes = primetable<int>(sqN);
    sum_lo.resize(sqN + 1);
    sum_hi.resize(sqN + 1);
    FOR3(i, 1, sqN + 1) sum_lo[i] = f(i) - 1;
    FOR3(i, 1, sqN + 1) sum_hi[i] = f(double(N) / i) - 1;
    for (int p: primes) {
      ll pp = ll(p) * p;
      if (pp > N) break;
      int R = min(sqN, N / pp);
      int M = sqN / p;
      T x = sum_lo[p - 1];
      T fp = sum_lo[p] - sum_lo[p - 1];
      for (int i = 1; i <= M; ++i) sum_hi[i] -= fp * (sum_hi[i * p] - x);
      for (int i = M + 1; i <= R; ++i) sum_hi[i] -= fp * (sum_lo[N / (double(i) * p)] - x);
      for (int n = sqN; n >= pp; --n) sum_lo[n] -= fp * (sum_lo[n / p] - x);
    }
    calculated = 1;
  }

  void calc_count() {
    calc([](ll x) -> T { return x; });
  }

  void calc_sum() {
    calc([](ll x) -> T {
      ll a = x, b = x + 1;
      if (!(x & 1)) a /= 2;
      if (x & 1) b /= 2;
      return T(a) * T(b);
    });
  }
};
#line 3 "nt/count_by_factor_type.hpp"

// factor type: 降順 270 -> (3,1,1)
// N=10^9: 1324 種類, 0.4sec
// https://atcoder.jp/contests/xmascon20/tasks/xmascon20_d
map<vc<int>, ll> count_by_factor_type(ll N) {
  ll sqN = sqrtl(N);
  auto P = primetable<int>(sqN);
  PrimeSum<ll> X(N);
  X.calc_count();

  // 1 and prime
  map<vc<int>, ll> ANS;
  ANS[vc<int>()] = 1;
  if (X[N] > 0) ANS[vc<int>({1})] = X[N];

  auto add = [&](vc<int> F, int k) -> vc<int> {
    int p = len(F);
    F.eb(k);
    while (0 < p && F[p - 1] < F[p]) { swap(F[p - 1], F[p]), --p; }
    return F;
  };

  // t = up_i^k のときに
  auto dfs = [&](auto& dfs, ll t, ll i, ll k, vc<int> U) -> void {
    // U * primes を追加する
    vc<int> nxt1 = add(U, k + 1);
    ANS[nxt1]++;
    vc<int> Uk = add(U, k);
    vc<int> nxt2 = add(Uk, 1);
    ll cnt = X[N / t] - X[P[i]];
    if (cnt > 0) ANS[nxt2] += X[N / t] - X[P[i]];
    ll lim = sqrtl(double(N) / t);
    if (P[i] <= lim) { dfs(dfs, t * P[i], i, k + 1, U); }
    FOR(j, i + 1, len(P)) {
      if (P[j] > lim) break;
      dfs(dfs, t * P[j], j, 1, Uk);
    }
  };
  FOR(i, len(P)) if (P[i] <= sqN) dfs(dfs, P[i], i, 1, {});
  return ANS;
}

Back to top page