\(g=\gcd(A,M)\) とし、 \(B=B_1g+B_2\) \((0 \le B_2 < g)\) とすると \(\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)}\) が成り立ちます。したがって、 \(g>1\) の場合も \(g=1\) の場合に帰着させることができます。以降は \(g=1\) の場合を考えます。

\(N=N_1M+N_2\) \((0\le N_2 < M)\) とすると \(\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}\) が成り立つため、 \(N\geq M\) の場合も \(N<M\) に帰着させることができます。以降は \(N<M\) も仮定します。

これらの仮定の元では \( (Ak+B)\bmod M\) \((k=0,1,\ldots,N-1)\) の値は相異なることに留意してください。

以下では Sqrt Heuristic for Floor Sums という記事で扱われているテクニックを用いた解法を紹介します。

この問題を「集合 \(\lbrace (Ak+B)\bmod M\ |\ 0\le k < N\rbrace\) をいくつかの等差数列の互いに素な和集合に分解し、それらで独立した問題を考える」という方針で解くことを考えます。いくつかの等差数列の互いに素な和集合に分解する方針として、順に \(2\) つの方針を考えてみます。

1.

集合 \(\lbrace (Ak+B)\bmod M\ |\ 0\le k < N\rbrace\) を \(\displaystyle \left\lfloor\frac{Ak+B}M \right\rfloor\) の値によって分別することを考えます。例えば \((N,M,A,B)=(4,5,2,1)\) の場合 \((Ak+B) \bmod M\) の値は \(k\) の昇順に \(1,3,0,2\) ですが、これを \((1,3)\) と \((0,2)\) に分解するということです。

このように分解すると、商が同じ要素をまとめたことから分解された各数列は全て等差数列となることが分かります。したがって、この方針では \(\displaystyle O\left(\frac{AN}{M}\right)\) 個の等差数列に分解することができます。

2.

1.で考えた方針を元により少ない個数に分解することを考えます。

\(D=O(\sqrt{M})\) となる定数 \(D\) を取ります。もし \(N\le D\) なら長さ \(1\) の \(N\) 個の等差数列に分けることにします。この場合の個数は \(O(N)=O(\sqrt{M})\) 個です。以降は \(N >D\) の場合を考えます。

まず、 \((Ak+B) \bmod M\) の最初 \(D\) 項を調べます。これら \(D\) 個の値から \(2\) つ異なる要素をとってきた時の差の最小値は鳩の巣原理より \(\displaystyle \frac{M}{D}\) 以下になることが分かります。最小値を満たす \(2\) つの要素のインデックスの差を \(d\) \((\le D)\) 、要素の差を \(h\) \(\displaystyle\left(\le\frac{M}{D}\right)\) とします。

数列 \((Ak+B) \bmod M\) を \(k \bmod d\) の値によって分別することを考えます。分別後は \((Adk+B') \bmod M\) と表される長さ \(\displaystyle O\left(\frac{N}{d} \right)\) の \(d\) 個の数列に分解されます。ここで、 \(Ad \equiv \pm h \bmod M\) が成り立つため、分解後の各数列は 1. と同じ手法を用いることで \(\displaystyle O\left(\frac{h\times \frac{N}d}M\right)=O\left(\frac{hN}{dM}\right)\) 個の等差数列に分解できます。つまり、全体として \(\displaystyle O\left(\frac{hN}{M}\right)=O\left(\frac{N}{D} \right)=O(\sqrt{M})\) 個の等差数列に分解できます。

以上より、 \(N\le D\) の場合も \(N>D\) の場合も \(O(\sqrt{M})\) 個の等差数列に分解できることが分かります。

あとは各等差数列に対して計算し、その値の総和を計算すれば良いです。各等差数列に対して等比数列の総和の計算が \(O(\log M)\) 時間かかるので計算量は全体として \(O(\sqrt{M}\log M)\) となります。

以上を適切に実装することでこの問題に正答することができます。