类欧几里得算法

我们要快速求以下函数的值：

$$f(a,b,c,n)=\sum_{i=0}^n \lfloor \frac{ai+b}{c} \rfloor$$

$$g(a,b,c,n)=\sum_{i=0}^n i \lfloor \frac{ai+b}{c} \rfloor$$

$$h(a,b,c,n)=\sum_{i=0}^n \lfloor \frac{ai+b}{c} \rfloor^2$$

f

先考虑 $a\ge c$ 或 $b\ge c$ 的情况。我们考虑将 $a$ 和 $b$ 对 $c$ 取模。

$$\begin{aligned} f(a,b,c,n)&=\sum_{i=0}^n \lfloor \frac{ai+b}{c} \rfloor\\ &=\sum_{i=0}^n\lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor+i\lfloor \frac{a}{c} \rfloor+\lfloor \frac{b}{c} \rfloor\\ &=(\sum_{i=0}^n \lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor)+(\sum_{i=0}^ni\lfloor \frac{a}{c} \rfloor)+(\sum_{i=0}^n\lfloor \frac{b}{c} \rfloor)\\ &=f(a\bmod c,b\bmod c,c,n)+\frac{n(n+1)}{2}\lfloor \frac{a}{c} \rfloor+(n+1)\lfloor \frac{b}{c} \rfloor \end{aligned}$$

此时 $a$ 变为 $a\bmod c$，$b$ 变为 $b\bmod c$。我们接下来考虑 $a<c,b<c$ 的情况。设 $m=\lfloor \frac{an+b}{c} \rfloor$，则

$$\begin{aligned} f(a,b,c,n)&=\sum_{i=0}^n \lfloor \frac{ai+b}{c} \rfloor\\ &=\sum_{i=0}^n \sum_{j=0}^{\lfloor \frac{ai+b}{c} \rfloor-1} 1\\ &=\sum_{i=0}^n \sum_{j=0}^{m-1} [j<\lfloor \frac{ai+b}{c} \rfloor]\\ &=\sum_{j=0}^{m-1} \sum_{i=0}^n [j<\lfloor \frac{ai+b}{c} \rfloor]\\ \end{aligned}$$

我们推导 $[j<\lfloor \frac{ai+b}{c} \rfloor]$ 的值：

$$\begin{aligned} j<\lfloor \frac{ai+b}{c} \rfloor & \iff j+1\le\lfloor \frac{ai+b}{c} \rfloor \\ &\iff j+1\le \frac{ai+b}{c} \\ &\iff ai\ge cj+c-b \\ &\iff ai> cj+c-b-1 \\ &\iff i> \lfloor\frac{cj+c-b-1}{a}\rfloor \end{aligned}$$

所以 $[j<\lfloor \frac{ai+b}{c} \rfloor]=[i> \lfloor\frac{cj+c-b-1}{a}\rfloor]$。我们继续推导：

$$\begin{aligned} f(a,b,c,n)&=\sum_{j=0}^{m-1} \sum_{i=0}^n [j<\lfloor \frac{ai+b}{c} \rfloor]\\ &=\sum_{j=0}^{m-1} \sum_{i=0}^n [i> \lfloor\frac{cj+c-b-1}{a}\rfloor]\\ &=\sum_{j=0}^{m-1} n-\lfloor\frac{cj+c-b-1}{a}\rfloor\\ &=nm-\sum_{j=0}^{m-1} \lfloor\frac{cj+c-b-1}{a}\rfloor\\ &=nm-f(c,c-b-1,a,m-1) \end{aligned}$$

最后，边界情况是 $a=0$。显然此时 $f(a,b,c,n)=(n+1)\lfloor \frac{b}{c} \rfloor$。综上：

$$f(a,b,c,n)=\begin{cases} f(a\bmod c,b\bmod c,c,n)+\frac{n(n+1)}{2}\lfloor \frac{a}{c} \rfloor+(n+1)\lfloor \frac{b}{c} \rfloor & a\ge c 或 b\ge c \\ nm-f(c,c-b-1,a,m-1) & a<c,b<c \\ (n+1)\lfloor \frac{b}{c} \rfloor & a=0 \\ \end{cases}$$

我们观察式子中 $a,c$ 位置上的变化。在第一个式子中，$a$ 对 $c$ 取模。在第二个式子中，$a,c$ 交换了位置。这样交替进行，直到 $a=0$ 的过程，与欧几里得算法中辗转相除的过程相同。因此直接递归的时间复杂度是 $O(\log n)$。

g

接下来的推推导过程与 $f$ 类似，但是复杂了很多。

先考虑 $a\ge c$ 或 $b\ge c$ 的情况

$$\begin{aligned} g(a,b,c,n)&=\sum_{i=0}^n i \lfloor \frac{ai+b}{c} \rfloor\\ &=\sum_{i=0}^n i(\lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor+i\lfloor \frac{a}{c} \rfloor+\lfloor \frac{b}{c} \rfloor)\\ &=(\sum_{i=0}^n i\lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor)+(\sum_{i=0}^ni^2\lfloor \frac{a}{c} \rfloor)+(\sum_{i=0}^n i\lfloor \frac{b}{c} \rfloor)\\ &=g(a\bmod c,b\bmod c,c,n)+\frac{n(n+1)(2n+1)}{6}\lfloor \frac{a}{c} \rfloor+\frac{n(n+1)}{2}\lfloor \frac{b}{c} \rfloor \end{aligned}$$

然后考虑 $a<c,b<c$ 的情况。省略了一些在推导 $f$ 时推导过的东西。

$$\begin{aligned} g(a,b,c,n)&=\sum_{i=0}^n i \lfloor \frac{ai+b}{c} \rfloor\\ &=\sum_{j=0}^{m-1} \sum_{i=0}^n i\times[i> \lfloor\frac{cj+c-b-1}{a}\rfloor]\\ &=\sum_{j=0}^{m-1} \sum_{i=\lfloor\frac{cj+c-b-1}{a}\rfloor+1}^n i\\ \end{aligned}$$

我们设 $t=\lfloor\frac{cj+c-b-1}{a}\rfloor$。继续推导

$$\begin{aligned} g(a,b,c,n)&=\sum_{j=0}^{m-1} \sum_{i=t+1}^n i\\ &=\sum_{j=0}^{m-1}\frac{(t+1+n)(n-t)}{2}\\ &=\sum_{j=0}^{m-1}\frac{(n^2+n)-(t^2+t)}{2}\\ &=\frac{1}{2}\sum_{j=0}^{m-1}(n^2+n)-(t^2+t)\\ &=\frac{1}{2}[nm(n+1)-\sum_{j=0}^{m-1}\lfloor\frac{cj+c-b-1}{a}\rfloor^2-\sum_{j=0}^{m-1}\lfloor\frac{cj+c-b-1}{a}\rfloor]\\ &=\frac{1}{2}[nm(n+1)-h(c,c-b-1,a,m-1)-f(c,c-b-1,a,m-1)]\\ \end{aligned}$$

最后考虑边界情况 $a=0$。显然此时 $g(a,b,c,n)=\frac{n(n+1)}{2}\lfloor \frac{b}{c} \rfloor$。综上：

$$g(a,b,c,n)=\begin{cases} g(a\bmod c,b\bmod c,c,n)+\frac{n(n+1)(2n+1)}{6}\lfloor \frac{a}{c} \rfloor+\frac{n(n+1)}{2}\lfloor \frac{b}{c} \rfloor & a\ge c 或 b\ge c \\ \frac{1}{2}[nm(n+1)-h(c,c-b-1,a,m-1)-f(c,c-b-1,a,m-1)] & a<c,b<c \\ \frac{n(n+1)}{2}\lfloor \frac{b}{c} \rfloor & a=0 \\ \end{cases}$$

h

先考虑 $a\ge c$ 或 $b\ge c$ 的情况

$$\begin{aligned} h(a,b,c,n)&=\sum_{i=0}^n \lfloor \frac{ai+b}{c} \rfloor^2\\ &=\sum_{i=0}^n (\lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor+i\lfloor \frac{a}{c} \rfloor+\lfloor \frac{b}{c} \rfloor)^2\\ &=\sum_{i=0}^n \lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor^2+(i\lfloor \frac{a}{c} \rfloor)^2+\lfloor \frac{b}{c} \rfloor^2+2i\lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor\lfloor \frac{a}{c} \rfloor+2\lfloor \frac{(a \bmod c)i+b\bmod c}{c} \rfloor\lfloor \frac{b}{c} \rfloor+2i\lfloor \frac{a}{c} \rfloor\lfloor \frac{b}{c} \rfloor\\ &=h(a\bmod c,b \bmod c,c,n)+\frac{n(n+1)(2n+1)}{6}\lfloor \frac{a}{c} \rfloor^2+(n+1)\lfloor \frac{b}{c} \rfloor^2+2g(a\bmod c,b \bmod c,c,n)\lfloor \frac{a}{c} \rfloor+2f(a\bmod c,b \bmod c,c,n) \lfloor \frac{b}{c} \rfloor+n(n+1)\lfloor \frac{a}{c} \rfloor\lfloor \frac{b}{c} \rfloor \end{aligned}$$

整理得：

$$h(a,b,c,n)=h(a\bmod c,b \bmod c,c,n)+2f(a\bmod c,b \bmod c,c,n) \lfloor \frac{b}{c} \rfloor+2g(a\bmod c,b \bmod c,c,n)\lfloor \frac{a}{c} \rfloor+\frac{n(n+1)(2n+1)}{6}\lfloor \frac{a}{c} \rfloor^2+(n+1)\lfloor \frac{b}{c} \rfloor^2+n(n+1)\lfloor \frac{a}{c} \rfloor\lfloor \frac{b}{c} \rfloor$$

然后考虑 $a<c,b<c$ 的情况。平方不太好处理，我们考虑平方分成两个部分。

$$n^2=2\frac{n(n+1)}{2}-n=2\sum_{i=1}^ni-n$$

所以：

$$\begin{aligned} h(a,b,c,n)&=\sum_{i=0}^n \lfloor \frac{ai+b}{c} \rfloor^2\\ &=\sum_{i=0}^n\left[\left(2\sum_{j=1}^{\lfloor \frac{ai+b}{c} \rfloor}j\right)-\lfloor \frac{ai+b}{c} \rfloor^2\right]\\ &=\left(2\sum_{i=0}^n\sum_{j=1}^{\lfloor \frac{ai+b}{c} \rfloor}j\right)-f(a,b,c,n) \end{aligned}$$

我们接着推导式子的前半部分：

$$\begin{aligned} \sum_{i=0}^n\sum_{j=1}^{\lfloor \frac{ai+b}{c} \rfloor}j&=\sum_{i=0}^n\sum_{j=0}^{\lfloor \frac{ai+b}{c} \rfloor-1}j+1\\ &=\sum_{j=0}^{m-1}(j+1)\sum_{i=0}^n[j<\lfloor \frac{ai+b}{c} \rfloor]\\ &=\sum_{j=0}^{m-1}(j+1)\sum_{i=0}^n[i> \lfloor\frac{cj+c-b-1}{a}\rfloor]\\ &=\sum_{j=0}^{m-1}(j+1)(n-\lfloor\frac{cj+c-b-1}{a}\rfloor)\\ &=\sum_{j=0}^{m-1}jn+n-\lfloor\frac{cj+c-b-1}{a}\rfloor-j\lfloor\frac{cj+c-b-1}{a}\rfloor\\ &=\frac{nm(m-1)}{2}+nm-f(c,c-b-1,a,m-1)-g(c,c-b-1,a,m-1)\\ &=\frac{nm(m+1)}{2}-f(c,c-b-1,a,m-1)-g(c,c-b-1,a,m-1)\\ \end{aligned}$$

带入到原式，得

$$\begin{aligned} h(a,b,c,n)&=\left(2\sum_{i=0}^n\sum_{j=1}^{\lfloor \frac{ai+b}{c} \rfloor}j\right)-f(a,b,c,n)\\ &=nm(m+1)-2f(c,c-b-1,a,m-1)-2g(c,c-b-1,a,m-1)-f(a,b,c,n) \end{aligned}$$

最后考虑边界情况 $a=0$。显然此时 $h(a,b,c,n)=(n+1)\lfloor \frac{b}{c} \rfloor^2$。综上：

$$h(a,b,c,n)=\begin{cases} h(a\bmod c,b \bmod c,c,n)+2f(a\bmod c,b \bmod c,c,n) \lfloor \frac{b}{c} \rfloor+2g(a\bmod c,b \bmod c,c,n)\lfloor \frac{a}{c} \rfloor+\frac{n(n+1)(2n+1)}{6}\lfloor \frac{a}{c} \rfloor^2+(n+1)\lfloor \frac{b}{c} \rfloor^2+n(n+1)\lfloor \frac{a}{c} \rfloor\lfloor \frac{b}{c} \rfloor & a\ge c 或 b\ge c \\ nm(m+1)-2f(c,c-b-1,a,m-1)-2g(c,c-b-1,a,m-1)-f(a,b,c,n) & a<c,b<c \\ (n+1)\lfloor \frac{b}{c} \rfloor^2 & a=0 \\ \end{cases}$$

实现

如果 $f,g,h$ 分开实现，时间复杂度会无法保证。例如，在求 $g(a,b,c,n)$ 时，如果满足 $a<c,b<c$，那么程序会同时调用 $f$ 和 $h$，导致递归形成很多分支，时间复杂度也就无法保证。

所以我们将 $f,g,h$ 放在同一个函数实现。代码没有任何难度，抄上面的公式即可。公式很复杂，注意不要抄错了。

const int mod = 998244353, inv2 = 499122177, inv6 = 166374059;
struct S
{
    ll f, g, h;
    S(ll f, ll g, ll h) : f(f), g(g), h(h) {}
};
S solve(ll a, ll b, ll c, ll n)
{
    if (a == 0) return S((n + 1) * (b / c) % mod, n * (n + 1) / 2 % mod * (b / c) % mod, (n + 1) * (b / c) % mod * (b / c) % mod);
    if (a >= c || b >= c)
    {
        S x = solve(a % c, b % c, c, n);
        ll f = (x.f + n * (n + 1) / 2 % mod * (a / c) % mod + (n + 1) * (b / c) % mod) % mod;
        ll g = (x.g + (n * (n + 1) % mod * (2 * n + 1) % mod * inv6 % mod) * (a / c) % mod + n * (n + 1) / 2 % mod * (b / c) % mod) % mod;
        ll h = (x.h + 2 * x.f % mod * (b / c) % mod + 2 * x.g % mod * (a / c) % mod + (n * (n + 1) % mod * (2 * n + 1) % mod * inv6 % mod) * (a / c) % mod * (a / c) % mod + (n + 1) * (b / c) % mod * (b / c) % mod + n * (n + 1) % mod * (a / c) % mod * (b / c) % mod) % mod;
        return S(f, g, h);
    }
    else
    {
        ll m = (a * n + b) / c;
        S x = solve(c, c - b - 1, a, m - 1);
        ll f = (n * m % mod - x.f + mod) % mod;
        ll g = ((m * n % mod * (n + 1) % mod - x.h - x.f) % mod + mod) % mod * inv2 % mod;
        ll h = ((n * m % mod * (m + 1) % mod - 2 * x.f - 2 * x.g - f) % mod + mod) % mod;
        return S(f, g, h);
    }
}