Sum of Floor of Linear
(math/number-theory/floor_sum.hpp)

• View this file on GitHub
• Last update: 2024-01-08 17:31:43+09:00
• Include: #include "math/number-theory/floor_sum.hpp"

Description

一次関数の床関数の和 $\sum_{i=0}^{N-1} \left\lfloor \frac{Ai + B}{M} \right\rfloor$ を再帰的に計算する．

計算量はユークリッドの互除法と同様に対数時間となるが，どの変数に依存するのかはよくわかっていない (僕が理解していない)．

• long long floor_sum(long long n, long long m, long long a, long long b)
  • $\sum_{i=0}^{N-1} \left\lfloor \frac{Ai + B}{M} \right\rfloor$ を計算する

Verified with

• test/yosupo/sum_of_floor_of_linear.test.cpp

Code

#pragma once

long long floor_sum(long long n, long long m, long long a, long long b) {
    long long sum = 0;
    if (a >= m) {
        sum += (a / m) * n * (n - 1) / 2;
        a %= m;
    }
    if (b >= m) {
        sum += (b / m) * n;
        b %= m;
    }
    long long y = (a * n + b) / m;
    if (y == 0) return sum;
    long long x = (m * y - b + a - 1) / a;
    sum += (n - x) * y + floor_sum(y, a, m, a * x - m * y + b);
    return sum;
}

#line 2 "math/number-theory/floor_sum.hpp"

long long floor_sum(long long n, long long m, long long a, long long b) {
    long long sum = 0;
    if (a >= m) {
        sum += (a / m) * n * (n - 1) / 2;
        a %= m;
    }
    if (b >= m) {
        sum += (b / m) * n;
        b %= m;
    }
    long long y = (a * n + b) / m;
    if (y == 0) return sum;
    long long x = (m * y - b + a - 1) / a;
    sum += (n - x) * y + floor_sum(y, a, m, a * x - m * y + b);
    return sum;
}

Back to top page