【动态规划】【SDOI2017】序列计数

【动态规划】【SDOI2017】序列计数

题目描述

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

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

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

$1\le n\le {10}^9,1\le m\le 2\times {10}^7,1\le p\le 100$。

算法概述

先不考虑质数的限制，对第一条限制做 $dp$ ，容易得到：

设 $f_{i,j}$ 为前 $i$ 个数，模 $p$ 的值为 $j$ 的方案数。

$$f_{i,j}=\sum _k{f}_{i-1,k}\times cnt{f}_{(j-k+p)\mathrm{\%}p}$$

其中 $cnt{f}_i$ 为 $1\to m$ 中模 $p$ 为 $i$ 的数字个数。

同理，考虑质数限制，设 $g_{i,j}$ 为前 $i$ 个数，只取合数，模 $p$ 的值为 $j$ 的方案数。

$$g_{i,j}=\sum _k{g}_{i-1,k}\times cnt{g}_{(j-k+p)\mathrm{\%}p}$$

其中 $cnt{g}_i$ 为 $1\to m$ 的合数中模 $p$ 为 $i$ 的数字个数。

发现 $p$ 比较小，并且 $j=(k+(j-k+p))\mathrm{\%}p$ ，所以我们考虑矩阵乘法优化这个 $dp$ 。

考虑从 $i-1$ 转移到 $i$ 的时候，$(x+y)\mathrm{\%}p=z$ 的项都会转移到 $z$ 上，考虑这样设计矩阵：

$$\left[\begin{array}{ccccc}cnt{f}_0& cnt{f}_1& cnt{f}_2& \dots & cnt{f}_{p-1}\\ cnt{f}_{p-1}& cnt{f}_0& cnt{f}_1& \dots & cnt{f}_{p-2}\\ cnt{f}_{p-2}& cnt{f}_{p-1}& cnt{f}_0& \dots & cnt{f}_{p-3}\\ \dots & \dots & \dots & \dots & \dots \\ cnt{f}_1& cnt{f}_2& cnt{f}_3& \dots & cnt{f}_0\end{array}\right]$$

初始矩阵为

$$\left[\begin{array}{ccccc}cnt{f}_0& cnt{f}_1& cnt{f}_2& \dots & cnt{f}_{p-1}\end{array}\right]$$

如果不清楚为什么是这样，建议将两个矩阵手模一下，观察乘法过程。

$f_{n,0}$ 就是初始矩阵乘以转移矩阵的 $n-1$ 次方的第 $0$ 项。

$g_{n,0}$ 的求法同理。

根据补集转化，至少有一个质数的答案就是 $f_{n,0}-g_{n,0}$ 。

Code

#include<bits/stdc++.h>
using namespace std;
const int N = 2e6 + 5,MOD = 20170408;
int prime[N],cnt = 0,p,cntf[105],cntg[105],m,n;
typedef long long ll;
ll f[105],g[105];
bitset <N * 10> vis;
inline void init()
{
	cntf[1]++;cntg[1]++;
	for(int i = 2;i <= m;i++)
	{
		cntf[i % p]++;
		if(!vis[i]) prime[++cnt] = i;
		else cntg[i % p]++;
		for(int j = 1;j <= cnt && 1ll * i * prime[j] <= m;j++)
		{
			vis[i * prime[j]] = 1;
			if(!(i % prime[j])) break;
		}
	}
}
struct Matrix {
	ll a[105][105];
};
Matrix operator *(Matrix x,Matrix y)
{
	Matrix z;
	memset(z.a,0,sizeof(z.a));
	for(int i = 0;i <= p - 1;i++)
		for(int j = 0;j <= p - 1;j++)
			for(int k = 0;k <= p - 1;k++)
				z.a[i][j] = (z.a[i][j] + x.a[i][k] * y.a[k][j] % MOD) % MOD;
	return z;
}
inline Matrix ksm(Matrix base,int pts)
{
	Matrix ret;
	memset(ret.a,0,sizeof(ret.a));
	for(int i = 0;i <= p - 1;i++) ret.a[i][i] = 1;
	for(;pts > 0;pts >>= 1,base = base * base)
		if(pts & 1)
			ret = ret * base;
	return ret;
}
int main()
{
	cin>>n>>m>>p;
	init();
	Matrix kf,kg;
	for(int i = 0;i <= p - 1;i++)
		for(int j = 0;j <= p - 1;j++)
			kf.a[i][j] = cntf[(j - i + p) % p],kg.a[i][j] = cntg[(j - i + p) % p];
	kf = ksm(kf,n - 1);
	kg = ksm(kg,n - 1);
	for(int i = 0;i <= p - 1;i++) 
		for(int j = 0;j <= p - 1;j++)
			f[i] = (f[i] + cntf[j] * kf.a[j][i] % MOD) % MOD,g[i] = (g[i] + cntg[j] * kg.a[j][i] % MOD) % MOD;
	cout<<(f[0] - g[0] + MOD) % MOD;
	return 0;
}