Skip to content

floor_sum.hpp

SECTIONMath INCLUDEnoya/floor_sum.hpp

计算任意整数区间上的整除下取整之和，支持系数为负。

\[ \displaystyle F=\sum_{i=L}^{R-1}\left\lfloor\frac{ai+b}{m}\right\rfloor \]

Complexity: Time: O(log m). Space: O(1).

AC 记录：sum_of_floor_of_linear。

跳到代码 · GitHub ↗

Implementation

当前头文件，省略 include guard；依赖见 #include。

/// @complexity Time: O(log m).
/// Space: O(1).

#include <cassert>
#include <cstdint>
#include <utility>

namespace noya {

using floor_sum_result = __int128_t;

namespace floor_sum_internal {

inline floor_sum_result floor_div(floor_sum_result val, floor_sum_result mod) {
  floor_sum_result quo = val / mod;
  if (val % mod < 0) {
    quo--;
  }
  return quo;
}

} // namespace floor_sum_internal

/// @brief Sum floor((a*i+b)/mod) over i in [l, r), allowing signed
/// l, a, and b; the exact result must fit in signed 128 bits.
inline floor_sum_result floor_sum_range(std::int64_t l, std::int64_t r,
                                        std::int64_t a, std::int64_t b,
                                        std::int64_t mod) {
  assert(l <= r && mod > 0);
  using floor_sum_internal::floor_div;
  floor_sum_result n = floor_sum_result(r) - l;
  floor_sum_result m = mod;
  floor_sum_result k = a;
  floor_sum_result ntr = floor_sum_result(a) * l + b;
  floor_sum_result res = 0;

  floor_sum_result sq = floor_div(k, m);
  k -= sq * m;
  res += n * (n - 1) / 2 * sq;
  floor_sum_result iq = floor_div(ntr, m);
  ntr -= iq * m;
  res += n * iq;

  while (true) {
    if (k >= m) {
      res += n * (n - 1) / 2 * (k / m);
      k %= m;
    }
    if (ntr >= m) {
      res += n * (ntr / m);
      ntr %= m;
    }
    floor_sum_result mx = k * n + ntr;
    if (mx < m) {
      break;
    }
    n = mx / m;
    ntr = mx % m;
    std::swap(m, k);
  }
  return res;
}

} // namespace noya

#ifndef NOYA_FLOOR_SUM_HPP
#define NOYA_FLOOR_SUM_HPP 1

/// @complexity Time: O(log m).
/// Space: O(1).

#include <cassert>
#include <cstdint>
#include <utility>

namespace noya {

using floor_sum_result = __int128_t;

namespace floor_sum_internal {

inline floor_sum_result floor_div(floor_sum_result val, floor_sum_result mod) {
  floor_sum_result quo = val / mod;
  if (val % mod < 0) {
    quo--;
  }
  return quo;
}

} // namespace floor_sum_internal

/// @brief Sum floor((a*i+b)/mod) over i in [l, r), allowing signed
/// l, a, and b; the exact result must fit in signed 128 bits.
inline floor_sum_result floor_sum_range(std::int64_t l, std::int64_t r,
                                        std::int64_t a, std::int64_t b,
                                        std::int64_t mod) {
  assert(l <= r && mod > 0);
  using floor_sum_internal::floor_div;
  floor_sum_result n = floor_sum_result(r) - l;
  floor_sum_result m = mod;
  floor_sum_result k = a;
  floor_sum_result ntr = floor_sum_result(a) * l + b;
  floor_sum_result res = 0;

  floor_sum_result sq = floor_div(k, m);
  k -= sq * m;
  res += n * (n - 1) / 2 * sq;
  floor_sum_result iq = floor_div(ntr, m);
  ntr -= iq * m;
  res += n * iq;

  while (true) {
    if (k >= m) {
      res += n * (n - 1) / 2 * (k / m);
      k %= m;
    }
    if (ntr >= m) {
      res += n * (ntr / m);
      ntr %= m;
    }
    floor_sum_result mx = k * n + ntr;
    if (mx < m) {
      break;
    }
    n = mx / m;
    ntr = mx % m;
    std::swap(m, k);
  }
  return res;
}

} // namespace noya

#endif // NOYA_FLOOR_SUM_HPP

#include <cassert>
#include <cstdint>
#include <utility>

/// @complexity Time: O(log m).
/// Space: O(1).

namespace noya {

using floor_sum_result = __int128_t;

namespace floor_sum_internal {

inline floor_sum_result floor_div(floor_sum_result val, floor_sum_result mod) {
  floor_sum_result quo = val / mod;
  if (val % mod < 0) {
    quo--;
  }
  return quo;
}

} // namespace floor_sum_internal

/// @brief Sum floor((a*i+b)/mod) over i in [l, r), allowing signed
/// l, a, and b; the exact result must fit in signed 128 bits.
inline floor_sum_result floor_sum_range(std::int64_t l, std::int64_t r,
                                        std::int64_t a, std::int64_t b,
                                        std::int64_t mod) {
  assert(l <= r && mod > 0);
  using floor_sum_internal::floor_div;
  floor_sum_result n = floor_sum_result(r) - l;
  floor_sum_result m = mod;
  floor_sum_result k = a;
  floor_sum_result ntr = floor_sum_result(a) * l + b;
  floor_sum_result res = 0;

  floor_sum_result sq = floor_div(k, m);
  k -= sq * m;
  res += n * (n - 1) / 2 * sq;
  floor_sum_result iq = floor_div(ntr, m);
  ntr -= iq * m;
  res += n * iq;

  while (true) {
    if (k >= m) {
      res += n * (n - 1) / 2 * (k / m);
      k %= m;
    }
    if (ntr >= m) {
      res += n * (ntr / m);
      ntr %= m;
    }
    floor_sum_result mx = k * n + ntr;
    if (mx < m) {
      break;
    }
    n = mx / m;
    ntr = mx % m;
    std::swap(m, k);
  }
  return res;
}

} // namespace noya