Official

G - Sum of Pow of Mod of Linear Editorial by sounansya

公式解説の補足

公式解説では \(N < M\) に限定して考えていましたが、一般に \(\lbrace (k,(Ak+B) \bmod M)\ |\ 0\le k < N \rbrace\) を \(O(\sqrt N)\) 個の等差数列に分解することを考えます。

まず \(D=O(\sqrt N)\) となる定数 \(D\) を取り、\(\min(Ai \bmod M, (-Ai) \bmod M)\) が最小となる \(1\le i < D\) を探します。この \(i\) を \(i_0\) とします。

この \(i_0\) に対し、鳩の巣原理から \(\displaystyle \min(Ai_0 \bmod M, (-Ai_0) \bmod M) \le \frac{M}{D}\) が成り立つことに留意してください。

もし \(Ai_0 \bmod M > (-Ai_0) \bmod M\) が成り立つ場合は \((A,B) \leftarrow (-A,M-1-B)\) とした場合を考えることで \(\displaystyle Ai_0 \bmod M \le \frac{M}{D}\) となる場合に帰着させることができます。以降はこの場合を考えます。

\(\lbrace (k,(Ak+B) \bmod M)\ |\ 0\le k < N \rbrace\) を一旦 \(k \bmod i_0\) の値で分類します。各等差数列は \(d=0,1,\ldots,i_0-1\) に対し \(\displaystyle \left\lbrace (ki_0+d,(A(ki_0+d)+B) \bmod M)\ |\ 0\le k < \left\lceil \frac{N-d}{i_0}\right\rceil \right\rbrace\) と表されます。これを公式解説のステップ 1. のように前から愚直に等差数列となるギリギリまで取ると、\(\displaystyle Ai_0 \bmod M \le \frac{M}{D}\) より \(\displaystyle O\left(\frac{1}{D}\times \frac{N}{i_0}\right)\) 個の等差数列に分かれます。したがって、全体で \(\displaystyle O\left(\frac{N}{D}\right)=O(\sqrt{N})\) 個の等差数列に分けることができます。

実装例(C++)

#include <atcoder/modint>
#include <bits/stdc++.h>
using namespace std;
tuple<long, long, vector<tuple<long, long, long>>> sqrt_floor(long n, long m, long a, long b) {
	if (n == 0) return {0, 0, {}};
	a %= m, b %= m;
	const long D = sqrtl(n);
	long dx = 1, best = m;
	for (long i = 1; i <= D; i++) {
		const long v = a * i % m;
		const long val = min(v, m - v);
		if (best > val) best = val, dx = i;
	}
	long dy = a * dx % m;
	const bool flip = dy > m - dy;
	if (flip) {
		a = (m - a) % m;
		b = m - 1 - b;
		dy = m - dy;
	}
	vector<tuple<long, long, long>> ans;
	for (long di = 0; di < dx; di++) {
		const long y = (b + a * di) % m;
		const long cnt = (n - 1 - di) / dx + 1;
		if (dy == 0) {
			ans.emplace_back(di, y, cnt);
			continue;
		}
		const long lim = (dy * (cnt - 1) + y) / m;
		long le = 0;
		for (long k = 0; k <= lim; k++) {
			const long ri = k == lim ? cnt : (m * (k + 1) - y + dy - 1) / dy;
			ans.emplace_back(dx * le + di, (dy * le + y) % m, ri - le);
			le = ri;
		}
	}
	if (flip) {
		dy *= -1;
		for (auto &[x, y, len] : ans) y = m - 1 - y;
	}
	return {dx, dy, ans};
}
using modint = atcoder::modint;
modint mp(modint x, long n) { // 1 + x + x^2 + ... + x^{n-1}
	if (n == 0) return 0;
	if (n % 2) return 1 + x * mp(x, n - 1);
	return (1 + x) * mp(x * x, n / 2);
}
int main() {
	cin.tie(nullptr);
	ios::sync_with_stdio(false);
	int t;
	cin >> t;
	while (t--) {
		long n, m, a, b, x, r;
		cin >> n >> m >> a >> b >> x >> r;
		modint::set_mod(r);
		modint ans = 0;
		auto [_, yy, res] = sqrt_floor(n, m, a, b);
		for (auto [_, y, len] : res) {
			if (yy < 0) y += (len - 1) * yy;
			ans += modint(x).pow(y) * mp(modint(x).pow(abs(yy)), len);
		}
		cout << ans.val() << '\n';
	}
}