Skip to the content.

:heavy_check_mark: Floor Sum
(math/floor-sum.hpp)

使い方

以下の問題を解く.

整数 $n,m,a,b\ (m\neq 0)$ に対し以下の値を求めよ: \(f(n,m,a,b)=\sum_{i=0}^{n-1}\left\lfloor\frac{ai+b}{m}\right\rfloor\)

$an+b$ が unsigned long long におさまる必要がある.

アルゴリズム

再帰的に考える.

$n=0$ のとき $f(n,m,a,b)=0$ である.

簡単な変形によって $m\geq 1$ および $0\leq a\lt m, 0\leq b\lt m$ として構わない.

$a\geq 0$ のとき $an+b=mq+r\,(0\leq r\lt m)$ とする. $0\leq i\lt n$ に対し $0\leq\left\lfloor\frac{ai+b}{m}\right\rfloor\leq q$ であり,$\left\lfloor\frac{ai+b}{m}\right\rfloor\geq k$ となる $i$ の範囲は $\frac{mk-b}{a}\leq i$ であるから以下が成り立つ.

\[\begin{align*} f(n,m,a,b) &=\sum_{k=1}^{q}\left(n-\left\lceil\frac{mk-b}{a}\right\rceil\right)\\ &=\sum_{k=1}^{q}\left(n+\left\lfloor\frac{b-mk}{a}\right\rfloor\right)\\ &=\sum_{k=0}^{q-1}\left(n+\left\lfloor\frac{b-m(q-k)}{a}\right\rfloor\right)\\ &=\sum_{k=0}^{q-1}\left\lfloor\frac{mk+an+b-mq}{a}\right\rfloor\\ &=\sum_{k=0}^{q-1}\left\lfloor\frac{mk+r}{a}\right\rfloor\\ &=f(q,a,m,r) \end{align*}\]

$m,a$ について Euclid の互除法と同じ議論ができ,$O(\log m)$ 回の反復で停止することが示せる.

Verified with

Code

#pragma once

// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSumUnsigned(unsigned long long n, unsigned long long m, unsigned long long a, unsigned long long b) {
  assert(m != 0);
  if (n == 0) return 0;
  T res = 0;
  while (true) {
    if (a >= m) {
      unsigned long long q = a / m;
      res += T(q) * (n / 2) * ((n - 1) | 1);
      a -= m * q;
    }
    if (b >= m) {
      unsigned long long q = b / m;
      res += T(q) * n;
      b -= m * q;
    }
    unsigned long long y = a * n + b;
    if (y < m) break;
    n = y / m, b = y - m * n;
    swap(a, m);
  }
  return res;
}

// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSum(long long n, long long m, long long a, long long b) {
  assert(m != 0);
  if (n <= 0) return 0;
  if (m < 0) a = -a, b = -b, m = -m;
  T res = 0;
  if (a < 0) {
    long long q = (a + 1) / m - 1;
    res += T(q) * (n / 2) * ((n - 1) | 1);
    a -= m * q;
  }
  if (b < 0) {
    long long q = (b + 1) / m - 1;
    res += T(q) * n;
    b -= m * q;
  }
  return res + FloorSumUnsigned<T>(n, m, a, b);
}

/**
 * @brief Floor Sum
 * @docs docs/math/floor-sum.md
 */
#line 2 "math/floor-sum.hpp"

// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSumUnsigned(unsigned long long n, unsigned long long m, unsigned long long a, unsigned long long b) {
  assert(m != 0);
  if (n == 0) return 0;
  T res = 0;
  while (true) {
    if (a >= m) {
      unsigned long long q = a / m;
      res += T(q) * (n / 2) * ((n - 1) | 1);
      a -= m * q;
    }
    if (b >= m) {
      unsigned long long q = b / m;
      res += T(q) * n;
      b -= m * q;
    }
    unsigned long long y = a * n + b;
    if (y < m) break;
    n = y / m, b = y - m * n;
    swap(a, m);
  }
  return res;
}

// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSum(long long n, long long m, long long a, long long b) {
  assert(m != 0);
  if (n <= 0) return 0;
  if (m < 0) a = -a, b = -b, m = -m;
  T res = 0;
  if (a < 0) {
    long long q = (a + 1) / m - 1;
    res += T(q) * (n / 2) * ((n - 1) | 1);
    a -= m * q;
  }
  if (b < 0) {
    long long q = (b + 1) / m - 1;
    res += T(q) * n;
    b -= m * q;
  }
  return res + FloorSumUnsigned<T>(n, m, a, b);
}

/**
 * @brief Floor Sum
 * @docs docs/math/floor-sum.md
 */
Back to top page