bzoj 4818: [Sdoi2017]序列计数

题意：

Alice想要得到一个长度为n的序列，序列中的数都是不超过m的正整数，而且这n个数的和是p的倍数。Alice还希望这n个数中，至少有一个数是质数。Alice想知道，有多少个序列满足她的要求。

题解：

矩乘优化dp。
因为p很小，所以每个数直接$mod\ p$就好了。
$f[i][j]$表示前i个数，$mod\ p$为j的方案数。
显然可以矩乘。
然后容斥，将所有素数减去，再dp一次，两次答案相减就可以了。
code：

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<iostream>
#define LL long long
using namespace std;
const LL mod=20170408;
int n,m,p,h[110];
int prime[20000010],pr=0;
struct node{
    int n,m;
    LL a[110][110];
    node() {memset(a,0,sizeof(a));}
};
node operator * (node a,node b)
{
    node ans;ans.n=a.n;ans.m=b.m;
    for(int i=0;i<ans.n;i++)
        for(int j=0;j<ans.m;j++)
            for(int k=0;k<a.m;k++)
                ans.a[i][j]=(ans.a[i][j]+a.a[i][k]*b.a[k][j])%mod;
    return ans;
}
node operator ^ (node a,int b)
{
    node ans;ans.n=a.n;ans.m=a.m;
    for(int i=0;i<a.n;i++) ans.a[i][i]=1;
    while(b)
    {
        if(b&1) ans=ans*a;
        a=a*a;b>>=1;
    }
    return ans;
}
bool v[20000010];
void pre()
{
    memset(v,true,sizeof(v));
    for(int i=2;i<=20000000;i++)
    {
        if(v[i]) prime[++pr]=i;
        for(int j=1;j<=pr&&i*prime[j]<=20000000;j++)
        {
            v[i*prime[j]]=false;
            if(i%prime[j]==0) break;
        }
    }
}
int solve()
{
    node a;a.n=a.m=p;
    for(int i=0;i<p;i++)
        for(int j=0;j<p;j++)
            (a.a[i][(i+j)%p]+=h[j])%=mod;
    node f;f.n=1;f.m=p;f.a[0][0]=1;
    f=f*(a^n);
    return f.a[0][0];
}
int main()
{
    pre();
    scanf("%d %d %d",&n,&m,&p);
    for(int i=1;i<=m;i++) h[i%p]++;
    LL ans=solve();
    for(int j=1;j<=pr;j++)
    {
        if(prime[j]>m) break;
        h[prime[j]%p]--;
    }
    ans-=solve();
    printf("%lld",(ans+mod)%mod);
}