G - Sum of Prod of Mod of Linear Editorial
			
			by 
sounansya
			
		
		
		前提知識として一般化 floor sum が必要です。一般化 floor sum は、整数 \(p,q,N,M,A,B\) に対し \(\displaystyle f_{p,q}(N,M,A,B)=\sum_{i=0}^{N-1}i^p\left\lfloor \frac{Ai+B}M \right\rfloor^q\) を \(O((p+q)^4\log M)\) 時間で計算するアルゴリズムです。一般化 floor sum に関する詳しい説明が必要な方は こちらの記事 を参考にしてください。
対称性より \(B_1\le B_2\) を仮定します。
まず、求める式を式変形します。
\[ \begin{aligned} &\phantom{=} \sum_{k=0}^{N-1}\left\lbrace (Ak+B_1)\ \text{mod}\ M \right\rbrace\left\lbrace (Ak+B_2)\ \text{mod}\ M \right\rbrace\\ &=\sum_{k=0}^{N-1}\left((Ak+B_1)-M\left\lfloor \frac{Ak+B_1}M \right\rfloor \right)\left((Ak+B_2)-M\left\lfloor \frac{Ak+B_2}M \right\rfloor \right)\\ &=\sum_{k=0}^{N-1}(Ak+B_1)(Ak+B_2)-M\sum_{k=0}^{N-1}\left((Ak+B_1)\left\lfloor\frac{Ak+B_2}M \right\rfloor+(Ak+B_2)\left\lfloor\frac{Ak+B_1}M \right\rfloor \right)+M^2\sum_{k=0}^{N-1}\left\lfloor\frac{Ak+B_1}M \right\rfloor\left\lfloor\frac{Ak+B_2}M \right\rfloor \\ \end{aligned} \]
\(\displaystyle \sum_{k=0}^{N-1}(Ak+B_1)(Ak+B_2)-M\sum_{k=0}^{N-1}\left((Ak+B_1)\left\lfloor\frac{Ak+B_2}M \right\rfloor+(Ak+B_2)\left\lfloor\frac{Ak+B_1}M \right\rfloor \right)\) は一般化 floor sum を用いることで \(O(\log M)\) 時間で求めることができるので、結局 \(\displaystyle \sum_{k=0}^{N-1}\left\lfloor\frac{Ak+B_1}M \right\rfloor\left\lfloor\frac{Ak+B_2}M \right\rfloor\) を求めることができれば良いです。
\(\displaystyle c=\left\lfloor\frac{Ak+B_2}M \right\rfloor-\left\lfloor\frac{Ak+B_1}M \right\rfloor\) とすると、 \(c\) のとりうる値は \(0\) か \(1\) なので \(c(c-1)=0\) が成立します。したがって、
\[ \begin{aligned} &\phantom{=}c(c-1)\\ &=\left(\left\lfloor\frac{Ak+B_2}M \right\rfloor-\left\lfloor\frac{Ak+B_1}M \right\rfloor\right)\left(\left\lfloor\frac{Ak+B_2}M \right\rfloor-\left\lfloor\frac{Ak+B_1}M \right\rfloor-1\right)\\ &=\left\lfloor\frac{Ak+B_1}M \right\rfloor^2+\left\lfloor\frac{Ak+B_1}M \right\rfloor+\left\lfloor\frac{Ak+B_2}M \right\rfloor^2-\left\lfloor\frac{Ak+B_2}M \right\rfloor - 2\left\lfloor\frac{Ak+B_1}M \right\rfloor\left\lfloor\frac{Ak+B_2}M \right\rfloor\\ &=0 \end{aligned} \]
となります。ここから
\[\left\lfloor\frac{Ak+B_1}M \right\rfloor\left\lfloor\frac{Ak+B_2}M \right\rfloor=\frac12\left(\left\lfloor\frac{Ak+B_1}M \right\rfloor^2+\left\lfloor\frac{Ak+B_1}M \right\rfloor+\left\lfloor\frac{Ak+B_2}M \right\rfloor^2-\left\lfloor\frac{Ak+B_2}M \right\rfloor \right) \]
が分かるので、\(\displaystyle \sum_{k=0}^{N-1}\left\lfloor\frac{Ak+B_1}M \right\rfloor\left\lfloor\frac{Ak+B_2}M \right\rfloor\) も一般化 floor sum に帰着させることができます。
def floor_sum(n, m, a, b):
    if m == 0:
        return 0, 0, 0
    a1, a2 = a // m, a % m
    b1, b2 = b // m, b % m
    y = (a2 * n + b2) // m
    ff, gg, hh = floor_sum(y, a2, m, m + a2 - b2 - 1)
    nn = n * (n - 1) // 2
    f = n * y - ff
    g = n * y * y - ff - 2 * hh
    h = nn * y + (ff - gg) // 2
    g += n * (2 * n - 1) * (n - 1) * a1 * a1 // 6
    g += 2 * nn * a1 * b1
    g += b1 * b1 * n
    g += 2 * a1 * h
    g += 2 * b1 * f
    f += a1 * nn + b1 * n
    h += nn * (a1 * (2 * n - 1) + 3 * b1) // 3
    return f, g, h
import sys
input = sys.stdin.readline
for _ in range(int(input())):
    n, m, a, b1, b2 = map(int, input().split())
    if b1 >= b2:
        b1, b2 = b2, b1
    ans = 0
    ans += a * a * (n - 1) * n * (2 * n - 1) // 6
    ans += a * (b1 + b2) * (n - 1) * n // 2
    ans += b1 * b2 * n
    f1, g1, h1 = floor_sum(n, m, a, b1)
    f2, g2, h2 = floor_sum(n, m, a, b2)
    ans -= m * (a * h1 + b2 * f1)
    ans -= m * (a * h2 + b1 * f2)
    ans += m * m * (g1 + f1 + g2 - f2) // 2
    print(ans)
				posted:
				
				
				last update:
				
			
