[BZOJ 4818] [SDOI 2017] 序列计数

Description

Alice想要得到一个长度为 \(n\) 的序列，序列中的数都是不超过 \(m\) 的正整数，而且这 \(n\) 个数的和是 \(p\) 的倍数。

Alice还希望，这 \(n\) 个数中，至少有一个数是质数。

Alice想知道，有多少个序列满足她的要求。

Input

一行三个数，\(n,m,p\)。

Output

一行一个数，满足Alice的要求的序列数量，答案对 \(20170408\) 取模。

Sample Input

3 5 3

Sample Output

33

HINT

对 \(20\%\) 的数据，\(1\leq n,m\leq100\)
对 \(50\%\) 的数据，\(1\leq m \leq 100\)
对 \(80\%\) 的数据，\(1\leq m\leq 10^6\)
对 \(100\%\) 的数据，\(1\leq n \leq 10^9,1\leq m \leq 2\times 10^7,1\leq p\leq 100\)

Solution

设 \(f[i][j]\) 表示前 \(i\) 位之和模 \(p\) 等于 \(j\) 的方案数，\(g[i]\) 表示在模 \(p\) 意义下等于 \(i\) 的数的个数，有

\[f[i][j]=f[i-1][k]\cdot g[(j+p-k)\bmod p],0\le k < p \]

可构造出转移矩阵

\[\begin{bmatrix} g_0&g_{p-1}&\cdots&g_{1}\\ g_{1}&g_{0}&\cdots&g_{2}\\ \vdots&\vdots&\ddots&\vdots\\ g_{p-1}&g_{p-2}&\cdots&g_{0} \end{bmatrix} \begin{bmatrix} f_0\\ f_1\\ \vdots\\ f_{p-1} \end{bmatrix} \]

最后用总的答案减去不包含质数的答案。

再科普一下NTT做法：

\[(g_0+g_1x+g_2x^2+\cdots+g_{p-1}x^{p-1})^{n-1} \]

Code

#include <cstdio>
#include <cstring>

const int N = 20000002, mod = 20170408;
int n, m, p, pr[N], np[N], tot, f[102], g[102];

struct Matrix {
	int mat[102][102];
	Matrix(){ memset(mat, 0, sizeof mat); }
	Matrix operator * (const Matrix a) const {
		Matrix b;
		for (int i = 1; i <= p; ++i)
			for (int j = 1; j <= p; ++j)
				for (int k = 1; k <= p; ++k)
					b.mat[i][j] = (b.mat[i][j] + 1LL * a.mat[i][k] * mat[k][j]) % mod;
		return b;
	}
} unit, a, b;
void sieve() {
	np[1] = 1;
	for (int i = 2; i <= m; ++i) {
		if (!np[i]) pr[++tot] = i;
		for (int j = 1; j <= tot && i * pr[j] <= m; ++j) {
			np[i * pr[j]] = 1;
			if (i % pr[j] == 0) break;
		}
	}
}
int ksm(Matrix a, int b) {
	Matrix res = unit;
	for (; b; b >>= 1, a = a * a)
		if (b & 1) res = res * a;
	return res.mat[1][1];
}
int main() {
	scanf("%d%d%d", &n, &m, &p);
	sieve();
	for (int i = 1; i <= p; ++i) unit.mat[i][i] = 1;
	for (int i = 1; i <= m; ++i) ++f[i % p];
	for (int i = 1; i <= m; ++i) if (np[i]) ++g[i % p];
	for (int i = 1; i <= p; ++i) {
		int k = 1;
		for (int j = i - 1; j >= 0; --j, ++k) a.mat[i][k] = f[j], b.mat[i][k] = g[j];
		for (int j = p - 1; j >= i; --j, ++k) a.mat[i][k] = f[j], b.mat[i][k] = g[j];
	}
	printf("%d\n", (ksm(a, n) - ksm(b, n) + mod) % mod);
	return 0;
}