Skip to content

floor_sum.hpp

SECTIONMath INCLUDEnoya/floor_sum.hpp

Signed floor sums over an arbitrary half-open integer interval.

Verified by sum_of_floor_of_linear.

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

Implementation

View on GitHub

#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 value,
                                  floor_sum_result modulus) {
  floor_sum_result quotient = value / modulus;
  if (value % modulus < 0) {
    quotient--;
  }
  return quotient;
}

} // namespace floor_sum_internal

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

  floor_sum_result slope_quotient = floor_div(slope, m);
  slope -= slope_quotient * m;
  result += n * (n - 1) / 2 * slope_quotient;
  floor_sum_result intercept_quotient = floor_div(intercept, m);
  intercept -= intercept_quotient * m;
  result += n * intercept_quotient;

  while (true) {
    if (slope >= m) {
      result += n * (n - 1) / 2 * (slope / m);
      slope %= m;
    }
    if (intercept >= m) {
      result += n * (intercept / m);
      intercept %= m;
    }
    floor_sum_result maximum = slope * n + intercept;
    if (maximum < m) {
      break;
    }
    n = maximum / m;
    intercept = maximum % m;
    std::swap(m, slope);
  }
  return result;
}

} // 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 value,
                                  floor_sum_result modulus) {
  floor_sum_result quotient = value / modulus;
  if (value % modulus < 0) {
    quotient--;
  }
  return quotient;
}

} // namespace floor_sum_internal

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

  floor_sum_result slope_quotient = floor_div(slope, m);
  slope -= slope_quotient * m;
  result += n * (n - 1) / 2 * slope_quotient;
  floor_sum_result intercept_quotient = floor_div(intercept, m);
  intercept -= intercept_quotient * m;
  result += n * intercept_quotient;

  while (true) {
    if (slope >= m) {
      result += n * (n - 1) / 2 * (slope / m);
      slope %= m;
    }
    if (intercept >= m) {
      result += n * (intercept / m);
      intercept %= m;
    }
    floor_sum_result maximum = slope * n + intercept;
    if (maximum < m) {
      break;
    }
    n = maximum / m;
    intercept = maximum % m;
    std::swap(m, slope);
  }
  return result;
}

} // namespace noya