牛客小白月赛14 A 简单计数【矩阵快速幂】

链接：https://ac.nowcoder.com/acm/contest/879/A
来源：牛客网
 

题目描述

thewindandrainisinmywayandnevergoingawaythewindandrainisinmywayandnevergoingaway
你在一个有 n 个城市的国家中行走，城市从 1 到 n 依次编号
任意两个城市之间都有一条双向道路可以通行，且你可以花一天的时间从当前所在的城市到达任意一个别的城市
由于你比较闲的无聊，所以你不会连续两天都呆在同一个城市，也就是说每天你所在的城市都不相同（这句话的意思是，对于相邻的两天，你所在的城市应该不同）
一开始你在 1 号城市，求经过 k 天后你回到 1 号城市的方案数

当然如果不存在任意一种方案就输出 0 就好了

输入描述:

第一行两个整数 n,k

输出描述:

一行一个整数表示答案对 998244353 取模后的结果

示例1

输入

复制

1 1

输出

复制

0

备注:

maker.cpp如下：

#include "bits/stdc++.h"
using namespace std;

int main() {
srand((unsigned long long) new char);
int n = int(8e8) + rand() % int(1e8);
int k = int(8e8) + rand() % int(1e8);
printf("%d %d\n", n, k);
}

分析：设F(a,0)表示走a天回到1的方案数，F(a,1)表示走a天不回到1的方案数，容易得到：

整理一下得到：

构造矩阵

直接上矩阵快速幂即可。

#include "bits/stdc++.h"
using namespace std;
const int mod = 998244353;

struct Matrix
{
    int n;
    long long a[3][3];
    Matrix (int n):n(n)
    {
        for(int i=0;i<n;i++)
            for(int j=0;j<n;j++)a[i][j]=0;
    }

};
Matrix operator*(Matrix A,Matrix B)
{
    int n=A.n;
    Matrix res(n);
    for(int i=0;i<n;i++)
        for(int j=0;j<n;j++)
            for(int k=0;k<n;k++)
            {
                res.a[i][j]+=(long long)A.a[i][k]*B.a[k][j]%mod;
                res.a[i][j]%=mod;
            }
    return res;
}
Matrix operator^(Matrix A,int b)
{
    int n=A.n;
    Matrix res(n),base=A;
    for(int i=0;i<n;i++)res.a[i][i]=1;
    while(b)
    {
        if(b&1)res=res*base;
        base=base*base;
        b>>=1;
    }
    return res;
}

int main()
{
    long long n,k;
    cin>>n>>k;
    Matrix mat(2);
    mat.a[0][0]=n-1;
    mat.a[0][1]=n-2;
    mat.a[1][0]=(n-1)*(n-2)%mod;
    mat.a[1][1]=((n-2)*(n-2)%mod + n-1)%mod;
    mat=mat^(k/2);
    if(k&1)
    {
        long long ans=mat.a[0][1]*(n-1)%mod;
        cout<<ans<<endl;
    }
    else
    {
        long long ans=mat.a[0][0];
        cout<<ans<<endl;
    }
}