[ACM]【組合數計數/逆元】Atcoder 167 Colorful Blocks

Colorful Blocks

傳送門
題意：N個block來上色，M種顏色，要求上色方法數，滿足至多有K個相鄰pair同色。

思路：

我在求助題解之後才發現自己讀錯題了
題目的意思是pair，也就是說相鄰兩個同色。我一開始以爲是相鄰一串同色，最多不超過K個連續顏色。然後就不會啊（。）還以爲要弄圖論，連通塊，想來想去好像不能O(n)，O(n)的數學方法又想不到。
看了題解才知道，題目要求的是“至多K對相鄰兩個同色”。我們可以把這個條件轉化爲：除了最左邊那個以外的右邊的block中，至多K個與其左邊的block同色。
從0循環到K，i表示右邊blocks有i個與其相鄰左邊的block同色。也就是從n-1箇中選出i個，這i個顏色與左邊一致。最左邊的block的顏色有m種選法。右邊剩下的(n-i-1)個不能與其左邊的block顏色一致，因此只有(m-1)種選法。
$ans+=C{_{n-1}^{i}}\times m\times (m-1)^{n-1-i}$
過了第一關，剩下的就是如何計算組合數模。
我們注意到，此處是要求一連串的組合數模，因此使用線性法求連續逆元，再用公式法求組合數。（其他方法都會TLE）

代碼：

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
ll mod=998244353;
const ll maxn=200005;
ll fac[maxn],inv[maxn];
void inverse(int n){
	inv[1]=1;
	//線性求inv
	for(int i=2;i<=n;i++){
		inv[i]=(mod-mod/i)*inv[mod%i]%mod;
	}
}
//初始化，fac存儲factorial of x
//inv在inverse之後存儲x的inverse
//在for之後變爲存儲x!的inverse
void init(int n){
	fac[0]=inv[0]=1;//記得記錄初始位置
	inverse(n);
	for(int i=1;i<=n;i++){
		//這裏inv的遞推使用了inv爲積性函數的性質！
		fac[i]=fac[i-1]*i%mod,inv[i]=inv[i-1]*inv[i]%mod;
	}
}
ll c(int n,int m){
	//公式求c（逆元）
	return fac[n]*inv[m]%mod*inv[n-m]%mod;
}
ll quickp(ll a,ll n){
	ll ans=1;
	while(n){
		if(n&1) ans=ans*a%mod;
		a=a*a%mod;
		n>>=1;
	}
	return ans;
}
int main(){
	int n,m,k;
	scanf("%d%d%d",&n,&m,&k);
	init(n);
	ll p=quickp(m-1,n-1-k);
	ll ans=0;
	for(int i=k;i>=0;i--){
		//從k輪到0的目的是讓p(ie.(m-1)^{n-1-i})可以遞推
		//避免每次都快速冪
		ans=(ans+c(n-1,i)*m%mod*p%mod)%mod,p=p*(m-1)%mod;
	}
	printf("%lld\n",ans);
}