【题解】AT_abc279_g [ABC279G] At Most 2 Colors

设 \(f_{i,j}\) 表示考虑前 \(i\) 位，\(j\sim i\) 颜色相同，且 \(j-1\) 与 \(i\) 颜色不同，合法方案数。

对于 \(j<i\)，则 \(i\) 和 \(i-1\) 必须染相同色，只能从 \(f_{i-1,j}\) 转移；

对于 \(j=i\)，即 \(i\) 和 \(i-1\) 必须染不同色，所有 \(f_{i-1,j'}\) 都能产生贡献：

• 对于和 \(i\) 距离超过 \(m\) 的 \(j'\)，随便染和 \(i-1\) 不同的颜色都行，贡献为 \(f_{i-1,j'}\times (c-1)\)；

• 对于和 \(i\) 距离小于等于 \(m\) 的 \(j'\)，只能染成 \(j'-1\) 的颜色，贡献为 \(f_{i-1,j'}\)。

列出式子：

\[f_{i,j}= \begin{cases} f_{i,j-1},&j<i,\\ \left(\sum\limits_{j=1}^{i-m+1}f_{i-1,j'}\times (c-1)\right)+\left(\sum\limits_{j=i-m+2}^{i-1} f_{i-1,j'}\right),&j=i \end{cases} \]

直接按照这个式子 dp 是 \(O(n^2)\) 的可以发现对于 \(j<i\) 的情况直接继承 \(i-1\) 的答案，每次只要更新 \(j=i\) 的情况，所以第一维直接扔掉；\(j=i\) 可以使用前缀和计算 \(f\) 一段区间的和，每次转移只要 \(O(1)\)，故总时间复杂度 \(O(n)\)。

#include <bits/stdc++.h>
#define mem(a,v) memset(a,v,sizeof(a))
#define endl '\n'
#define FILE(x) freopen(x".in","r",stdin),freopen(x".out","w",stdout);
#define pii pair<int,int>
#define pll pair<long long,long long>
#define st first
#define nd second
#define pb push_back
using namespace std;
using ll=long long;
using lld=long double;
const int N=5e6+10;
const ll mod=998244353;
ll n,k,c;
ll f[N],g[N]; // g 是前缀和数组
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0),cout.tie(0);
    cin>>n>>k>>c,f[1]=g[1]=c;
    for(ll i=2;i<=n;i++){
        ll l=max(i-k+1,1ll);
        f[i]=(g[i-1]-g[l]+g[l]*(ll)(c-1)%mod)%mod,
        (f[i]+=mod)%=mod,
        g[i]=(g[i-1]+f[i])%mod;
    }
    cout<<g[n];
    return 0;
}