mod/range_freq_of_linear.hpp

• View this file on GitHub
• Last update: 2023-11-10 22:48:13+09:00
• Include: #include "mod/range_freq_of_linear.hpp"

Depends on

• mod/floor_sum_of_linear.hpp

Verified with

• test/library_checker/math/min_of_mod_of_linear_2.test.cpp
• test/mytest/range_freq_of_linear.test.cpp
• test/yukicoder/2280.test.cpp

Code

#include "mod/floor_sum_of_linear.hpp"

// L <= x < R のうちで、(ax+b mod) in [lo, hi) となるものの個数
ll range_freq_of_linear(ll L, ll R, ll a, ll b, ll mod, ll lo, ll hi) {
  if (lo >= hi) return 0;
  assert(0 <= lo && lo < hi && hi <= mod);

  i128 x1 = floor_sum_of_linear(L, R, a, b - lo, mod);
  i128 x2 = floor_sum_of_linear(L, R, a, b - hi, mod);
  return x1 - x2;
}

#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 2 "mod/range_freq_of_linear.hpp"

// L <= x < R のうちで、(ax+b mod) in [lo, hi) となるものの個数
ll range_freq_of_linear(ll L, ll R, ll a, ll b, ll mod, ll lo, ll hi) {
  if (lo >= hi) return 0;
  assert(0 <= lo && lo < hi && hi <= mod);

  i128 x1 = floor_sum_of_linear(L, R, a, b - lo, mod);
  i128 x2 = floor_sum_of_linear(L, R, a, b - hi, mod);
  return x1 - x2;
}

Back to top page