nt/range_rational_count.hpp

• View this file on GitHub
• Last update: 2024-09-14 09:20:23+09:00
• Include: #include "nt/range_rational_count.hpp"

Depends on

• enumerate/floor_range.hpp
• mod/floor_sum_of_linear.hpp
• nt/array_on_floor.hpp
• nt/mertens.hpp
• nt/mobius_table.hpp
• nt/primetable.hpp
• nt/zeta.hpp

Verified with

• test/3_yukicoder/2266.test.cpp

Code

#include "nt/mertens.hpp"
#include "mod/floor_sum_of_linear.hpp"

// 最大分母 N を指定する
// 既約分数を数えたり k 番目を求めたりする
struct Range_Rational_Count {
  u32 N;
  u64 total;
  Mertens<int> M;
  Range_Rational_Count(u32 N) : N(N), M(N) { total = count(1, 1); }

  // [0, a/b)
  u64 count(u32 a, u32 b) {
    assert(a <= b);
    if (a == 0) return 0;
    // [0,a/b]
    u64 ans = 1;
    floor_range(N, [&](u32 q, u32 l, u32 r) -> void {
      ans += u64(M[q]) * floor_sum_of_linear<u64, u64>(l, r, a, 0, b);
    });
    // a/b
    if (b <= N) --ans;
    return ans;
  }

  // [0,1) の中で k 番目
  pair<u32, u32> kth(u64 k) {
    assert(k < total);
    u32 int_part = k / total;
    k %= total;
    map<pair<u32, u32>, u64> MP;
    auto query = [&](u32 a, u32 b) -> u64 {
      pair<u32, u32> k = {a, b};
      if (MP.count(k)) return MP[k];
      return MP[k] = count(a, b);
    };
    // k 個以下なものの max
    u32 a = 0, b = 1, c = 1, d = 1;
    while (b + d <= N) {
      // 右に進む
      u32 l = 0, r = 1;
      while (b + r * d <= N && query(a + r * c, b + r * d) <= k) {
        l = r, r = 2 * r;
      }
      while (l + 1 < r) {
        u32 m = l + (r - l) / 2;
        (query(a + m * c, b + m * d) <= k ? l : r) = m;
      }
      a += l * c, b += l * d;
      // 左に進む
      l = 0, r = 1;
      while (r * b + d <= N && query(r * a + c, r * b + d) > k) {
        l = r, r = 2 * r;
      }
      while (l + 1 < r) {
        u32 m = l + (r - l) / 2;
        (query(m * a + c, m * b + d) > k ? l : r) = m;
      }
      c += l * a, d += l * b;
    }
    return {int_part * b + a, b};
  }
};

#line 1 "nt/range_rational_count.hpp"

#line 1 "nt/array_on_floor.hpp"
// N=10 だと dat = {dp[1], dp[2], dp[3], dp[5], dp[10]} みたいになる
// hashmap より数倍高速
template <typename T>
struct Array_On_Floor {
  u64 N;
  u32 t, sq;
  vc<T> dat;
  Array_On_Floor() {}
  Array_On_Floor(u64 N, T default_value = T{}) : N(N) {
    assert(N <= u64(1) << 50);
    sq = sqrtl(N);
    t = (u64(sq) * sq + sq <= N ? sq : sq - 1);
    dat.resize(t + sq, default_value);
  }

  u32 size() { return dat.size(); }

  T& operator[](u64 d) {
    int i = get_index(d);
    return dat[i];
  }

  inline u32 get_index(u64 d) {
    assert(d > 0);
    if (d <= t) return d - 1;
    return dat.size() - u32(double(N) / d);
  }

  // dat[i] に対応する floor
  u64 get_floor(u32 i) { return (i < t ? 1 + i : double(N) / (t + sq - i)); }

  template <typename F>
  void enumerate_all(F f) {
    FOR(i, len(dat)) { f(get_floor(i), dat[i]); }
  }
};
#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/zeta.hpp"

template <typename T>
void divisor_zeta(vc<T>& A) {
  assert(A[0] == 0);
  int N = len(A) - 1;
  auto P = primetable(N);
  for (auto&& p: P) { FOR3(x, 1, N / p + 1) A[p * x] += A[x]; }
}

template <typename T>
void divisor_mobius(vc<T>& A) {
  assert(A[0] == 0);
  int N = len(A) - 1;
  auto P = primetable(N);
  for (auto&& p: P) { FOR3_R(x, 1, N / p + 1) A[p * x] -= A[x]; }
}

template <typename T>
void multiplier_zeta(vc<T>& A) {
  assert(A[0] == 0);
  int N = len(A) - 1;
  auto P = primetable(N);
  for (auto&& p: P) { FOR3_R(x, 1, N / p + 1) A[x] += A[p * x]; }
}

template <typename T>
void multiplier_mobius(vc<T>& A) {
  assert(A[0] == 0);
  int N = len(A) - 1;
  auto P = primetable(N);
  for (auto&& p: P) { FOR3(x, 1, N / p + 1) A[x] -= A[p * x]; }
}
#line 2 "nt/mobius_table.hpp"

template<typename T>
vc<T> mobius_table(int N){
  vc<T> mu(N + 1);
  mu[1] = T(1);
  divisor_mobius(mu);
  return mu;
}
#line 1 "enumerate/floor_range.hpp"
// 商が q の区間 [l,r) を q について昇順

template <typename F>
void floor_range(u64 N, F f, bool Q_ASC = true, bool INCLUDE_Q_IS_0 = false) {
  u64 sq = sqrtl(N);
  u32 n = (sq * sq + sq <= N ? sq : sq - 1);
  if (Q_ASC) {
    if (INCLUDE_Q_IS_0) f(0, N + 1, infty<ll>);
    for (u32 q = 1; q <= n; ++q) { f(q, N / (q + 1) + 1, N / q + 1); }
    for (u32 l = sq; l >= 1; --l) { f(N / l, l, l + 1); }
  } else {
    for (u32 l = 1; l <= sq; ++l) { f(N / l, l, l + 1); }
    for (u32 q = n; q >= 1; --q) { f(q, N / (q + 1) + 1, N / q + 1); }
    if (INCLUDE_Q_IS_0) f(0, N + 1, infty<ll>);
  }
}
#line 4 "nt/mertens.hpp"

template <typename T>
struct Mertens {
  Array_On_Floor<T> sum;
  Mertens() {}
  Mertens(u64 N, u64 K = -1) { build(N, K); }
  void build(u64 N, u64 K = -1) {
    sum = Array_On_Floor<T>(N);
    if (K == u64(-1)) { K = pow(N, 0.67); }
    vc<T> A = mobius_table<T>(K);
    FOR(k, 1, K) A[k + 1] += A[k];
    FOR(i, len(sum)) {
      u64 n = sum.get_floor(i);
      if (n <= K) {
        sum.dat[i] = A[n];
        continue;
      }
      T ans = 1;
      floor_range(n, [&](u64 q, u64 l, u64 r) -> void {
        if (q == n) return;
        ans -= sum[q] * T(r - l);
      });
      sum.dat[i] = ans;
    }
  }
  T operator[](u64 n) { return sum[n]; }
};
#line 2 "mod/floor_sum_of_linear.hpp"

// sum_{x in [L,R)} floor(ax + b, mod)
// I は範囲内で ax+b がオーバーフローしない程度
template <typename O = i128, typename I = long long>
O floor_sum_of_linear(I L, I R, I a, I b, I mod) {
  assert(L <= R);
  O res = 0;
  b += L * a;
  I N = R - L;

  if (b < 0) {
    I k = ceil(-b, mod);
    b += k * mod;
    res -= O(N) * O(k);
  }

  while (N) {
    I q;
    tie(q, a) = divmod(a, mod);
    res += (N & 1 ? O(N) * O((N - 1) / 2) * O(q) : O(N / 2) * O(N - 1) * O(q));
    if (b >= mod) {
      tie(q, b) = divmod(b, mod);
      res += O(N) * q;
    }
    tie(N, b) = divmod(a * N + b, mod);
    tie(a, mod) = mp(mod, a);
  }
  return res;
}
#line 4 "nt/range_rational_count.hpp"

// 最大分母 N を指定する
// 既約分数を数えたり k 番目を求めたりする
struct Range_Rational_Count {
  u32 N;
  u64 total;
  Mertens<int> M;
  Range_Rational_Count(u32 N) : N(N), M(N) { total = count(1, 1); }

  // [0, a/b)
  u64 count(u32 a, u32 b) {
    assert(a <= b);
    if (a == 0) return 0;
    // [0,a/b]
    u64 ans = 1;
    floor_range(N, [&](u32 q, u32 l, u32 r) -> void {
      ans += u64(M[q]) * floor_sum_of_linear<u64, u64>(l, r, a, 0, b);
    });
    // a/b
    if (b <= N) --ans;
    return ans;
  }

  // [0,1) の中で k 番目
  pair<u32, u32> kth(u64 k) {
    assert(k < total);
    u32 int_part = k / total;
    k %= total;
    map<pair<u32, u32>, u64> MP;
    auto query = [&](u32 a, u32 b) -> u64 {
      pair<u32, u32> k = {a, b};
      if (MP.count(k)) return MP[k];
      return MP[k] = count(a, b);
    };
    // k 個以下なものの max
    u32 a = 0, b = 1, c = 1, d = 1;
    while (b + d <= N) {
      // 右に進む
      u32 l = 0, r = 1;
      while (b + r * d <= N && query(a + r * c, b + r * d) <= k) {
        l = r, r = 2 * r;
      }
      while (l + 1 < r) {
        u32 m = l + (r - l) / 2;
        (query(a + m * c, b + m * d) <= k ? l : r) = m;
      }
      a += l * c, b += l * d;
      // 左に進む
      l = 0, r = 1;
      while (r * b + d <= N && query(r * a + c, r * b + d) > k) {
        l = r, r = 2 * r;
      }
      while (l + 1 < r) {
        u32 m = l + (r - l) / 2;
        (query(m * a + c, m * b + d) > k ? l : r) = m;
      }
      c += l * a, d += l * b;
    }
    return {int_part * b + a, b};
  }
};

Back to top page