G - Sum of Pow of Mod of Linear Editorial by en_translator
Let \(g=\gcd(A,M)\) and write \(B=B_1g+B_2\) \((0 \le B_2 < g)\). Then \(\displaystyle X^{(Ak+B) \bmod M} = X^{B_2}{X}^{g \left(\left(\frac{A}{g}k+B_1\right)\bmod \frac{M}{g}\right)}\), so cases where \(g>1\) can be reduced to \(g=1\). From now on, we assume \(g=1\).
If we write \(N=N_1M+N_2\) \((0\le N_2 < M)\), we have \(\displaystyle \sum_{k=0}^{N-1}X^{(Ak+B)\bmod M}=N_1\sum_{k=0}^{M-1}X^k+\sum_{k=0}^{N_2-1}X^{(Ak+B)\bmod M}\), so cases where \(N\geq M\) can be boiled down to \(N<M\). Thus we will also assume \(N<M\).
Under this assumption, note that the values \( (Ak+B)\bmod M\) \((k=0,1,\ldots,N-1)\) are pairwise distinct.
We will introduce a solution that exploits the trick explained in the article Sqrt Heuristic for Floor Sums.
We aim at solving this problem by the following approach: express the set \(\lbrace (Ak+B)\bmod M\ |\ 0\le k < N\rbrace\) as the union of disjoint sets, each consisting of integers that form an arithmetic progression, and solve the problem individually for each of them. To decompose the set into disjoint arithmetic progressions, we consider the following two approaches:
1.
Consider classifying the set \(\lbrace (Ak+B)\bmod M\ |\ 0\le k < N\rbrace\) by the values of \(\displaystyle \left\lfloor\frac{Ak+B}M \right\rfloor\). For example, if \((N,M,A,B)=(4,5,2,1)\), the values of \((Ak+B) \bmod M\) are \(1,3,0,2\) in ascending order of \(k\), which we decompose into \((1,3)\) and \((0,2)\).
By this decomposition, each group with the same quotient forms an arithmetic progression. Thus, this approach allows us to decompose it into \(\displaystyle O\left(\frac{AN}{M}\right)\) arithmetic progressions.
2.
Consider refining approach 1 to reduce the number of groups.
Take a constant \(D\) with \(D=O(\sqrt{M})\). If \(N\le D\), we split it into \(N\) arithmetic progressions, each of length \(1\). In this case, the number of groups is \(O(N)=O(\sqrt{M})\). Hereinafter, we assume \(N >D\).
First, inspect the first \(D\) terms of \((Ak+B) \bmod M\). By the pigeonhole principle, the minimum difference between any two different values among them does not exceed \(\displaystyle \frac{M}{D}\). Let \(d\) \((\le D)\) be the difference between the indices of the two elements that achieve the minimum, and \(h\) \(\displaystyle\left(\le\frac{M}{D}\right)\) be the difference of the elements.
Consider decomposing the sequence \((Ak+B) \bmod M\) by the values \(k \bmod d\), so as to obtain \(d\) sequences of length \(\displaystyle O\left(\frac{N}{d} \right)\), each represented as \((Adk+B') \bmod M\). Here, since we have \(Ad \equiv \pm h \bmod M\), each of the decomposed sequence can be further split into \(\displaystyle O\left(\frac{h\times \frac{N}d}M\right)=O\left(\frac{hN}{dM}\right)\) arithmetic progressions by applying the same method as approach 1. In total, the original set can be decomposed into \(\displaystyle O\left(\frac{hN}{M}\right)=O\left(\frac{N}{D} \right)=O(\sqrt{M})\) arithmetic progressions.
Hence, we see that we can obtain a decomposition into \(O(\sqrt{M})\) arithmetic progressions, both when \(N\le D\) and \(N>D\).
All that left is to evaluate the expression for each arithmetic progression, and find the sum of the results. For each group, evaluating the sum of the geometric progression costs \(O(\log M)\) time, so the total complexity is \(O(\sqrt{M}\log M)\).
The problem can be solved by appropriately implementing the problem.
posted:
last update: