P8229 [AGM 2022 资格赛] 抛硬币

\(\mathcal Link\)

直接推式子。

枚举第一次硬币反面朝下的位置。

\[\begin{aligned} E(n)&=\sum_{i=0}^{n-1} p^i(1-p)(k^i+E(n-i-1)) \\ &=\sum_{i=0}^{n-1} (pk)^i(1-p)+\sum_{i=0}^{n-1} p^i(1-p)E(n-i-1)\\ \end{aligned}\]

令 \(f(n)=\sum\limits_{i=0}^{n-1} (pk)^i(1-p),g(n)=\sum\limits_{i=0}^{n-1} p^i(1-p)E(n-i-1)\)

可得：

\[\begin{aligned}f(n)&=p\cdot k\cdot f(n-1)+(1-p)\\ g(n)&=(1-p)f(n-1)+g(n-1)\end{aligned}\]

然后这个式子可以矩阵加速。复杂度 \(\mathcal O(\log n)\)。可能需要卡常。

#include <cstdio>
#include <cctype>
#include <cstring>
#include <algorithm>
using namespace std;
#define getchar getchar_unlocked
struct Ios{}io;
template <typename _tp>
Ios &operator >>(Ios &in,_tp &x){
    x=0;int w=0;char c=getchar();
    for(;!isdigit(c);w|=c=='-',c=getchar());
    for(;isdigit(c);x=x*10+(c^'0'),c=getchar());
    if(w) x=-x;
    return in;
}
typedef long long ll;
const int mod=998244353;
struct Matrix{
    int a[4][4];
    void clear(){
        a[1][1]=0;a[1][2]=0;a[1][3]=0;
        a[2][1]=0;a[2][2]=0;a[2][3]=0;
        a[3][1]=0;a[3][2]=0;a[3][3]=0;
    }
    void init(){
        a[1][1]=1;a[1][2]=0;a[1][3]=0;
        a[2][1]=0;a[2][2]=1;a[2][3]=0;
        a[3][1]=0;a[3][2]=0;a[3][3]=1;
    }
};
Matrix operator *(Matrix a,Matrix b){
    Matrix ans;ans.clear();
    ans.a[1][1]=(1ll*a.a[1][1]*b.a[1][1]+1ll*a.a[1][2]*b.a[2][1]+1ll*a.a[1][3]*b.a[3][1])%mod;
    ans.a[1][2]=(1ll*a.a[1][1]*b.a[1][2]+1ll*a.a[1][2]*b.a[2][2]+1ll*a.a[1][3]*b.a[3][2])%mod;
    ans.a[1][3]=(1ll*a.a[1][1]*b.a[1][3]+1ll*a.a[1][2]*b.a[2][3]+1ll*a.a[1][3]*b.a[3][3])%mod;
    ans.a[2][1]=(1ll*a.a[2][1]*b.a[1][1]+1ll*a.a[2][2]*b.a[2][1]+1ll*a.a[2][3]*b.a[3][1])%mod;
    ans.a[2][2]=(1ll*a.a[2][1]*b.a[1][2]+1ll*a.a[2][2]*b.a[2][2]+1ll*a.a[2][3]*b.a[3][2])%mod;
    ans.a[2][3]=(1ll*a.a[2][1]*b.a[1][3]+1ll*a.a[2][2]*b.a[2][3]+1ll*a.a[2][3]*b.a[3][3])%mod;
    ans.a[3][1]=(1ll*a.a[3][1]*b.a[1][1]+1ll*a.a[3][2]*b.a[2][1]+1ll*a.a[3][3]*b.a[3][1])%mod;
    ans.a[3][2]=(1ll*a.a[3][1]*b.a[1][2]+1ll*a.a[3][2]*b.a[2][2]+1ll*a.a[3][3]*b.a[3][2])%mod;
    ans.a[3][3]=(1ll*a.a[3][1]*b.a[1][3]+1ll*a.a[3][2]*b.a[2][3]+1ll*a.a[3][3]*b.a[3][3])%mod;
    return ans;
}
Matrix ans;
int main(){
    int T;io>>T;
    while(T--){
        ll n,k,p;io>>n>>k>>p;
        ans.init();
        Matrix tmp;tmp.clear();
        tmp.a[1][1]=p*k%mod;tmp.a[1][2]=(1-p+mod)%mod;
        tmp.a[2][2]=1;
        tmp.a[3][1]=(1-p+mod)%mod;tmp.a[3][3]=1;
        while(n){
            if(n&1) ans=ans*tmp;
            n>>=1;tmp=tmp*tmp;
        }
        Matrix x;x.clear();
        x.a[1][3]=1;
        ans=x*ans;
        printf("%d\n",(ans.a[1][1]+ans.a[1][2])%mod);
    }
    return 0;
}