mod/mod_sum_of_linear.hpp

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

Depends on

• mod/floor_sum_of_linear.hpp

Verified with

• test/mytest/mod_sum_of_linear.test.cpp

Code

#include "mod/floor_sum_of_linear.hpp"

i128 mod_sum_of_linear(ll L, ll R, ll a, ll b, ll mod) {
  /*
  sum_{x in [L,R)} (ax + b mod mod)
  */
  i128 s = a * L + b;
  i128 t = a * (R - 1) + b;
  i128 sum = (s + t) * (R - L) / 2;
  i128 fsum = floor_sum_of_linear(L, R, a, b, mod);
  return sum - fsum * mod;
}

#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/mod_sum_of_linear.hpp"

i128 mod_sum_of_linear(ll L, ll R, ll a, ll b, ll mod) {
  /*
  sum_{x in [L,R)} (ax + b mod mod)
  */
  i128 s = a * L + b;
  i128 t = a * (R - 1) + b;
  i128 sum = (s + t) * (R - L) / 2;
  i128 fsum = floor_sum_of_linear(L, R, a, b, mod);
  return sum - fsum * mod;
}

Back to top page