AtCoder ABC #279 (トヨタシステムズプログラミングコンテスト2022) : G - At Most 2 Colors

こちらは短く書けてびっくりした。
https://atcoder.jp/contests/abc279/tasks/abc279_g

問題

N個のマスが1列に並んでおり、それぞれC色のいずれかで塗る。
連続するKマスに3色以上現れないような塗り方は何通りか。

解法

f(n,m) := 左端からnマスまで塗ったとき、nマス目と異なる色のマスが出てくる最も右の位置がmであるような塗り方
とする。nマス全部1色の場合は除く。
そうすると解は以下の和となる。

• 2色以上使う場合、sum(f(N,*))通り。
• 1色だけで塗る場合、C通り。

あとはf(n,m)の遷移を考える。

• nマス目とn+1マス目を同じ色で塗る場合を考えると、f(n+1,m) += f(n,m)
• nマス目まで1色で塗り、n+1マス目を違う色で塗る場合を考えると、f(n+1,n) += C*(C-1)
• m≦n-Kの場合、nマス目とn+1マス目を違う色で塗る場合を考えるとf(n+1,n) += sum(f(n,m))*C

f(n,*)をインラインで更新することを考えると、結局nごとに下2つの分だけ加算すればいいことになる。
以下ではBITを使っているが、普通に累積和を使って解くこともできそう。

ll N,K,C;
const ll mo=998244353;

template<class V, int ME> class BIT_mod {
public:
	V bit[1<<ME];
	BIT_mod(){ZERO(bit);};
	V operator()(int e) { if(e<0) return 0; ll s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s%mo;}
	void add(int e,V v) { e++; v=(v%mo+mo)%mo; while(e<=1<<ME) { bit[e-1]+=v; bit[e-1] -= (bit[e-1]>=mo)?mo:0; e+=e&-e;}}
};
BIT_mod<ll,21> bt;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>C;
	
	for(i=2;i<=N;i++) {
		//新規色は前の色がi-k前から
		ll nc=bt(i-K)*(C-2)%mo;
		//新たに2色作る場合も含め
		ll tot=(bt(i)+C*(C-1))%mo;
		bt.add(i-1,tot);
		bt.add(i-1,nc);
	}
	
	//1色
	ll ret=C+bt(N);
	cout<<ret%mo<<endl;
	
	
}

まとめ

Fよりだいぶさっくり解けて本番びっくりした。