Floor Sum

Returns the sum

\[ \sum_{i=0}^{n-1} \Bigl\lfloor \frac{a\times i + b}{m} \Bigr\rfloor \]

in \(O(\log m)\) complexity.

Code

My implementation was stolen straight from AtCoder Library.

Floor Sum

// Floor Sum {{{
// sum of floor[i=0...n-1]((a*i+b)/m)
unsigned long long __floor_sum_unsigned(unsigned long long n,
                                      unsigned long long m,
                                      unsigned long long a,
                                      unsigned long long b) {
  unsigned long long ans = 0;
  while (true) {
    if (a >= m) {
      ans += n * (n - 1) / 2 * (a / m);
      a %= m;
    }
    if (b >= m) {
      ans += n * (b / m);
      b %= m;
    }

    unsigned long long y_max = a * n + b;
    if (y_max < m) break;
    // y_max < m * (n + 1)
    // floor(y_max / m) <= n
    n = (unsigned long long)(y_max / m);
    b = (unsigned long long)(y_max % m);
    swap(m, a);
  }
  return ans;
}

long long floor_sum(long long n, long long m, long long a, long long b) {
  unsigned long long ans = 0;
  if (a < 0) {
    unsigned long long a2 = (a%m+m)%m;
    ans -= 1ULL * n * (n - 1) / 2 * ((a2 - a) / m);
    a = a2;
  }
  if (b < 0) {
    unsigned long long b2 = (a%m+m)%m;
    ans -= 1ULL * n * ((b2 - b) / m);
    b = b2;
  }
  return ans + __floor_sum_unsigned(n, m, a, b);
}
//}}}