bzoj 3328: PYXFIB 单位根反演

版权声明：2333 https://blog.csdn.net/liangzihao1/article/details/89467337

Description

Input
第一行一个正整数，表示数据组数据 ，接下来T行
每行三个正整数N,K,P

Output
T行，每行输出一个整数，表示结果

Sample Input
1
1 2 3

Sample Output
1

HINT

Source
By Wcmg

分析：
组合数很大，考虑化掉组合数。
化简上诉式子得到
$=\sum_{i=0}^{n}[i\ mod\ k==0]\binom{n}{i}*F(i)$
考虑单位根反演，得到
$=\sum_{i=0}^{n}\frac{1}{k}*\sum_{j=0}^{k-1}(w_{k}^j)^i*\binom{n}{i}*F(i)$
交换循环，
$=\frac{1}{k}*\sum_{j=0}^{k-1}\sum_{i=0}^{n}(w_{k}^j)^i*\binom{n}{i}*F(i)$
其中， $F(i)=A^i$， $A$是斐波那契数列的递推矩阵。后面就是一个二项式展开的形式。
$=\frac{1}{k}*\sum_{j=0}^{k-1}(w_{k}^jA+I)^n$
直接 $O(klogn)$算就好了。

代码：

/**************************************************************
    Problem: 3328
    User: liangzihao
    Language: C++
    Result: Accepted
    Time:8932 ms
    Memory:1300 kb
****************************************************************/
 
#include <iostream>
#include <cstdio>
#include <cmath>
#define LL long long
 
using namespace std;
 
LL G,n,w,wn,ans,mod;
int T,k,cnt;
int p[101];
 
struct matrix{
    LL a[3][3];
}A,B;
 
matrix operator *(matrix a,matrix b)
{
    matrix c;
    for (int i=0;i<=2;i++)
    {
        for (int j=0;j<=2;j++) c.a[i][j]=0;
    }
    for (int k=1;k<=2;k++)
    {
        for (int i=1;i<=2;i++)
        {
            for (int j=1;j<=2;j++) c.a[i][j]=(c.a[i][j]+a.a[i][k]*b.a[k][j]%mod)%mod;
        }
    }
    return c;
}
 
LL ksm(LL x,LL y)
{
    if (y==1) return x;
    LL c=ksm(x,y/2);
    c=c*c%mod;
    if (y&1) c=c*x%mod;
    return c;
}
 
void divide(int x)
{
    for (int i=2;i<=trunc(sqrt(x));i++)
    {
        if (x%i==0)
        {
            p[++cnt]=i;
            while (x%i==0) x/=i;
        }
    }
    if (x>1) p[++cnt]=x;
}
 
void findroot(int x)
{
    for (int i=2;i<x;i++)
    {
        int flag=0;
        for (int j=1;j<=cnt;j++)
        {
            if (ksm(i,(x-1)/p[j])==1)
            {
                flag=1;
                break;
            }
        }
        if (!flag)
        {
            G=i;
            return;
        }
    }
}
 
void power(LL p)
{
    if (p==1)
    {
        B=A;
        return;
    }
    power(p/2);
    B=B*B;
    if (p&1) B=B*A;
}
 
int main()
{
    scanf("%d",&T);
    while (T--)
    {
        scanf("%lld%d%lld",&n,&k,&mod);
        cnt=0;      
        divide(mod-1);
        findroot(mod);
        wn=ksm(G,(mod-1)/k);
        w=1;
        ans=0;
        for (int j=0;j<k;j++)
        {
            A.a[1][1]=1,A.a[1][2]=w;
            A.a[2][1]=w,A.a[2][2]=(w+1)%mod;
            power(n);
            ans=(ans+B.a[2][2])%mod;
            w=(w*wn)%mod;
        }
        ans=ans*ksm(k,mod-2)%mod;
        printf("%lld\n",ans);
    }
}