bzoj 4818 [Sdoi2017]序列计数 矩阵乘法优化dp+容斥

Description

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

一行三个数，n,m,p。
1<=n<=10^9,1<=m<=2×10^7,1<=p<=100
Output

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

3 5 3
Sample Output

33
HINT

传送门
还挺不错的题。
发现P很小，在不考虑是否是素数的情况下，
可以比较简单地想到一个dp：
令f[i][j]表示前i个位置，前i个位置上的数mod p的值为j的方案数。
那么方程很简单，枚举一个j的情况下，
$f[i][(j+k)\%p]+=f[i-1][k]*calc(j)，0<=j<p$
其中calc(x)表示满足i%p=x，1<=i<=m的i的个数，当然这个是很好算的。
那么时间复杂度O(N*$P^2$)，空间O(N*P)爆炸啦。。
但是发现f[i][x]只和f[i-1][y]有关，所以可以想到用矩阵乘法优化。
考虑构造一个P*1的答案矩阵和P*P的转移矩阵，
答案矩阵当然很明显啦：

⎡⎣⎢⎢⎢⎢f[i][0]f[i][1]⋮f[i][p−1]⎤⎦⎥⎥⎥⎥

$$\left[ \begin{matrix} f[i][0]\\ f[i][1]\\ \vdots \\ f[i][p-1]\\ \end{matrix} \right]$$
目标就是让这个矩阵乘上某个矩阵能够得到:
$$\left[ \begin{matrix} f[i+1][0]\\ f[i+1][1]\\ \vdots \\ f[i+1][p-1]\\ \end{matrix} \right]$$
因为根据dp式子，$f[i+1][(y+x)\%P]+=f[i][x]*calc(y)$
比如$f[i+1][1]=f[i][0]*calc(1)$
根据矩阵乘法，我们知道是一行乘一列，那么简单地对应对应就好了：
$$\left[ \begin{matrix} calc(0)&calc(P-1)&\cdots calc(1)\\ calc(1)&calc(0)&\cdots calc(2) \\ \vdots &\vdots&\vdots\\ calc(P-1)&calc(P-2)&\cdots calc(0)\\ \end{matrix} \right]$$
把这东西矩乘了，再乘上面那个记录f的矩阵，
那就能够在O($P^3*log(n))$的时间内做一次dp啦，
因为最终答案要求和是P的倍数，所以答案自然就是矩阵的[1][1]了


现在再来看看我们漏下了什么：质数的要求。
这个……似乎不是非常好弄？因为calc(i)不能分别出质数。
再看看dp里面能不能加入判断质数：显然不行了。。
那么考虑一下能不能容斥：求出和是P倍数的，但满足没有一个质数的方案数。
这个就很简单了，只要对于每个质数x，calc(x%P)减去1就好了。。
然后再进行一遍相同的过程，两个答案作差就是解。
所以正面考虑不出再去想想补集还是一个不错的方法。。


矩乘是第一个的列=第二个的行……
我一开始竟然搞错了= c =（没救）

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int 
    N=105,
    mod=20170408;
int n,m,P,pcnt;
int calc[105],prime[1300000];
bool notprime[20000000];
struct Matrix{
    int val[N][N],ln,lm;
    Matrix(){memset(val,0,sizeof(val));}
    Matrix operator *(Matrix x){
        Matrix c;
        for (int i=1;i<=ln;i++)
            for (int j=1;j<=x.lm;j++)
                for (int k=1;k<=lm;k++)
                    c.val[i][j]=(c.val[i][j]+(ll)val[i][k]*x.val[k][j]%mod)%mod;
        c.ln=ln,c.lm=x.lm;
        return c;
    }
}a,b;
Matrix ksm(Matrix a,int y){
    Matrix z;
    for (int i=1;i<=P;i++) z.val[i][i]=1;
    z.ln=z.lm=P;
    while (y){
        if (y&1) z=z*a;
        y>>=1,a=a*a;
    }
    return z;
}
void Get_Prime(){
    notprime[1]=1,pcnt=0;
    for (int i=2;i<=m;i++){
        if (!notprime[i]) prime[++pcnt]=i;
        for (int j=1;j<=pcnt;j++){
            if (prime[j]*i>m) break;
            notprime[prime[j]*i]=1;
            if (!(i%prime[j])) break;
        }
    }
}
void build_Matrix(){
    for (int i=1;i<=P;i++) b.val[i][1]=calc[i-1];
    b.ln=P,b.lm=1;
    a.ln=a.lm=P;
    for (int i=1;i<=P;i++){
        int tmp=i-1;
        for (int j=1;j<=P;j++){
            a.val[i][j]=calc[tmp];
            tmp--;
            if (tmp<0) tmp=P-1;
        }
    }
}
int main(){
    scanf("%d%d%d",&n,&m,&P);
    Get_Prime();
    for (int i=1;i<=m;i++) calc[i%P]++;
    build_Matrix();
    a=ksm(a,n-1)*b;
    int ans=a.val[1][1];

    for (int i=1;i<=pcnt;i++) calc[prime[i]%P]--;

    build_Matrix();
    a=ksm(a,n-1)*b;
    printf("%d\n",(ans-a.val[1][1]+mod)%mod);
    return 0;
}