BZOJ3328 PYZFIB 单位根反演、矩阵快速幂

传送门

---

看到下标为一个数的倍数的时候能够贡献答案以及\(p \equiv 1 \mod k\)，可以考虑单位根反演。

设\(T\)为斐波那契数列的转移矩阵\(=\left( \begin{array}{cccc} 1 & 1 \\ 1 & 0 \end{array} \right)\)，那么\(F_i = T^i[0][0]\)，下面的式子中直接使用\(T\)进行计算，最后取答案矩阵\(ans\)的\(ans[0][0]\)的值就是答案。

\(\begin{align*} \sum\limits_{i=0}^{\lfloor \frac{n}{k} \rfloor} T^{ik} \binom{n}{ik} &= \sum\limits_{i=0}^n T^i \binom{n}{i} [k \mid i] \\ &= \sum\limits_{i=0}^n T^i \binom{n}{i} \frac{1}{k} \sum\limits_{j=1}^k w_k^{ij} \\ &= \frac{1}{k} \sum\limits_{j=1}^k \sum\limits_{i=0}^n \binom{n}{i} (w_k^jT)^i \\ &= \frac{1}{k} \sum\limits_{j=1}^k (w_k^jT + E)^n \end{align*}\)

因为\(p \equiv 1 \mod k\)，所以\(w_k\)有意义，那么直接计算上式即可，复杂度\(O(Tklogn)\)。

#include<bits/stdc++.h>
//this code is written by Itst
using namespace std;

long long N , K , P;
struct matrix{
    int a[2][2];
    matrix(){memset(a , 0 , sizeof(a));}
    int* operator [](int x){return a[x];}

    matrix operator *(matrix b){
        matrix c;
        for(int i = 0 ; i < 2 ; ++i)
            for(int j = 0 ; j < 2 ; ++j)
                for(int k = 0 ; k < 2 ; ++k)
                    c[i][j] = (c[i][j] + 1ll * a[i][k] * b[k][j]) % P;
        return c;
    }

    matrix operator +(matrix b){
        matrix c;
        for(int i = 0 ; i < 2 ; ++i)
            for(int j = 0 ; j < 2 ; ++j)
                c[i][j] = (a[i][j] + b[i][j]) % P;
        return c;
    }

    matrix operator *(int x){
        matrix c;
        for(int i = 0 ; i < 2 ; ++i)
            for(int j = 0 ; j < 2 ; ++j)
                c[i][j] = 1ll * a[i][j] * x % P;
        return c;
    }
    
    matrix operator ^(long long x){
        matrix c , d = *this; c[0][0] = c[1][1] = 1;
        while(x){
            if(x & 1) c = c * d;
            d = d * d;
            x >>= 1;
        }
        return c;
    }
}T;

int poww(long long a , int b){
    int times = 1;
    while(b){
        if(b & 1) times = times * a % P;
        a = a * a % P;
        b >>= 1;
    }
    return times;
}

int findG(){
    vector < int > ins;
    int tmp = P - 1;
    for(int i = 2 ; i * i <= tmp ; ++i)
        if(tmp % i == 0){
            ins.push_back(P / i);
            while(tmp % i == 0) tmp /= i;
        }
    if(tmp - 1) ins.push_back(P / tmp);
    for(int i = 2 ; i < P ; ++i){
        bool flg = 1;
        for(int j = 0 ; j < ins.size() && flg ; ++j)
            flg &= poww(i , ins[j]) != 1;
        if(flg) return i;
    }
    return -1;
}

int main(){
    T[0][0] = T[0][1] = T[1][0] = 1;
    ios::sync_with_stdio(0);
    int t;
    for(cin >> t ; t ; --t){
        cin >> N >> K >> P;
        int W = poww(findG() , (P - 1) / K) , tms = W , ans = 0;
        for(int i = 1 ; i <= K ; ++i , tms = 1ll * tms * W % P){
            matrix tmp = T * tms;
            ++tmp[0][0]; ++tmp[1][1];
            ans = (ans + (tmp ^ N)[0][0]) % P;
        }
        cout << 1ll * ans * poww(K , P - 2) % P << endl;
    }
    return 0;
}