HDU 2865 Birthday Toy

同上题。中间选每种颜色都是等价的。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#define mod 1000000007
#define maxn 100050
using namespace std;
long long n,k,ans=0,tot=0,prime[maxn];
bool vis[maxn];
struct matrix
{
    long long a[4][4];
}base[32];
void get_table()
{
    for (long long i=2;i<=maxn-50;i++)
    {
        if (!vis[i]) {vis[i]=true;prime[++tot]=i;}
        for (long long j=1;j<=tot && i*prime[j]<=maxn-50;j++)
        {
            vis[i*prime[j]]=true;
            if (!i%prime[j]) break;
        }
    }
}
long long f_pow(long long x,long long y)
{
    x%=mod;
    long long ans=1,base=x;
    while (y)
    {
        if (y&1) ans=(ans*base)%mod;
        base=(base*base)%mod;
        y>>=1;
    }
    return ans;
}
long long inv(long long x) {return f_pow(x,mod-2);}
long long phi(long long x)
{
    long long ret1=1,ret2=1,top=x;
    for (register long long i=1;prime[i]*prime[i]<=top;i++)
    {
        if (x%prime[i]) continue;
        ret1*=(prime[i]-1);ret2*=prime[i];
        while (x!=1)
        {
            if (x%prime[i]) break;
            x/=prime[i];
        }
    }
    if (x!=1) {ret1*=(x-1);ret2*=x;}
    double ans=(double)top/ret2*ret1;
    return (long long)ans%mod;
}
matrix reset()
{
    matrix c;
    c.a[1][1]=k-3;c.a[1][2]=k-2;c.a[1][3]=0;
    c.a[2][1]=1;c.a[2][2]=0;c.a[2][3]=0;
    c.a[3][1]=0;c.a[3][2]=1;c.a[3][3]=0;
    return c;
}
matrix reset2()
{
    matrix c;
    for (long long i=1;i<=3;i++)
        for (long long j=1;j<=3;j++)
            c.a[i][j]=0;
    c.a[1][1]=(k-2)*(k-1)%mod*(k-3)%mod;
    c.a[2][1]=(k-2)*(k-1)%mod;
    c.a[3][1]=0;
    return c;
}
matrix I()
{
    matrix c;
    for (long long i=1;i<=3;i++)
        for (long long j=1;j<=3;j++)
            c.a[i][j]=(i==j);
    return c;
}
matrix operator * (matrix a,matrix b)
{
    matrix c;
    for (long long i=1;i<=3;i++)
        for (long long j=1;j<=3;j++)
            c.a[i][j]=0;
    for (long long k=1;k<=3;k++)
        for (long long i=1;i<=3;i++)
            for (long long j=1;j<=3;j++)
                c.a[i][j]=(c.a[i][j]+a.a[i][k]*b.a[k][j]%mod)%mod;
    return c;
}
long long f(long long x)
{
    if (x==1) return 0;
    else if (x==2) return (k-2)*(k-1)%mod;
    else if (x==3) return (k-2)*(k-1)%mod*(k-3)%mod;
    matrix ans=reset2();long long ret=0;x-=3;
    while (x)
    {
        if (x&1) ans=base[ret]*ans;
        ret++;x>>=1;
    }
    return ans.a[1][1]%mod;
}
void work()
{
    base[0]=reset();
    for (long long i=1;i<=31;i++) base[i]=base[i-1]*base[i-1];
    long long ans=0,top=sqrt(n);
    for (long long i=1;i<=top;i++)
    {
        if (n%i) continue;
        ans=(ans+f(i)*phi(n/i)%mod)%mod;
        if ((i!=top) || (top*top!=n)) ans=(ans+f(n/i)*phi(i)%mod)%mod;
    }
    ans=(ans*k)%mod;
    ans=(ans*inv(n))%mod;
    printf("%lld
",ans);
}
int main()
{
    get_table();
    while (scanf("%lld%lld",&n,&k)!=EOF)
        work();
    return 0;   
}