染色/CF1093F Vasya and Array 题解

· · 题解

题目：一个长度 n 的序列，每个元素有 1\sim m 的颜色，有些位置的颜色给定，其它的在 1\sim m 中任取，求将其染色的方案数，使得相邻 l 个元素颜色不全相同。

首先有一个 dp 方程， f_{i,x} 代表前 i 个染成最后一个是 x 的方案数，s_i=\sum_x f_{i,x}

当 [i-l+1,i] 可以被同时染成 x 时（也就是这个区间里面没有除了 x 之外的给定的颜色），

f_{i,x}=s_{i-1}-s_{i-l}+f_{i-l,x}

这个意义是，之前的随便选的方案数 s_i，减去这个位置恰好出现了 l 个连续的方案数 s_{i-l}-f_{i-l,x}，由于新加一个数只可能出现恰好 l 个的非法情况，所以这么减刚好就是真实的答案

否则，

f_{i,x}=s_{i-1}

因为随便选就一定合法。

直接维护上述式子是 \mathcal{O}(nm) 的

考虑复杂度瓶颈在于维护 f，但是其实并不需要这样，只需要维护 s 就好了。

考虑 f_{i,x}

我们发现，它有可能等于

f_{i,x}=s_{i-1}-s_{i-l}+f_{i-l,x}

但继续还可以拆

f_{i,x}=s_{i-1}-s_{i-l}+s_{i-l-1}-s_{i-2l}+f_{i-2l,x}

容易发现这就是一个这样的链条，直到某个位置停止，也就是说

f_{i,x}=s_{i-1}-s_{i-l}+s_{i-l-1}-s_{i-2l}+\cdots +s_{i-(k-1)l-1}-s_{i-kl}+s_{i-kl-1}

这个转移方程中没有 f。

首先，如果我们令 S_i=s_i+s_{i-l}+s_{i-2l}+\cdots，显然 S 可以直接维护，那么上面的转移方程显然可以变成

f_{i,x}=S_{i-1}-S_{i-(k+1)l-1}-(S_{i-l}-S_{i-(k+1)l})

那么现在问题在于如何求 k。

注意到对于一个 i 的不同 x，k 的取值最多只有2个不同的。

首先，每个 x 对应的 k 显然是每次向左走 l，第一次碰到不是 x 的颜色的位置的 k（根据之前转移方程的意义）

那么令每次走 l，第一次碰到的颜色是 c，位置是 k_1，第一次碰到不是 c 的位置是 k_2，注意这个 k_1,k_2，可以简单地扫一遍 \mathcal{O}(n) 求出来。

那么显然对于所有 x\neq c，对应的 k=k_1，否则若 x=c ，k=k_2

所以说，对于一个 i，有多少个 f_{i,x}=S_{i-1}-S_{i-(k_1+1)l-1}-(S_{i-l}-S_{i-(k_1+1)l})，多少个 f_{i,x}=S_{i-1}-S_{i-(k_2+1)l-1}-(S_{i-l}-S_{i-(k_2+1)l}) 都可以求出来，那么就可以 \mathcal{O}(1) 求出 s_i

所以说就可以 \mathcal{O}(n) 解决此题了！!1

#include<bits/stdc++.h>
#pragma GCC optimize("-Ofast")
using namespace std;
//dengyaotriangle!

const int maxn=100005*2;
constexpr int mdn=998244353;

struct mint{
    int v;
    inline operator int()const{return v;}
    inline mint(int v=0):v(v){}
    inline bool operator==(const mint&a)const{return v==a.v;}
    inline bool operator!=(const mint&a)const{return v!=a.v;}
    inline mint operator~ ()const{int r=1,b=v;for(int i=1;i<=(mdn-2);i<<=1){if((mdn-2)&i)r=r*(long long)b%mdn;b=b*(long long)b%mdn;}return mint(r);}
    inline mint operator- ()const{return mint(v?mdn-v:0);}
    inline mint operator+ (const mint&a)const{mint c(v+a.v);if(c.v>=mdn)c.v-=mdn;return c;}
    inline mint operator- (const mint&a)const{mint c(v-a.v);if(c.v<0)c.v+=mdn;return c;}
    inline mint operator* (const mint&a)const{return mint(v*(long long)a.v%mdn);}
    inline mint operator/ (const mint&a)const{return (*this)*~a;}
    inline mint&operator++(){v++;if(v==mdn)v=0;return*this;}
    inline mint&operator--(){v--;if(v==-1)v=mdn-1;return*this;}
    inline mint operator++(int){mint ret=*this;++*this;return ret;}
    inline mint operator--(int){mint ret=*this;--*this;return ret;}
    inline mint&operator+=(const mint&a){v+=a.v;if(v>=mdn)v-=mdn;return*this;}
    inline mint&operator-=(const mint&a){v-=a.v;if(v<0)v+=mdn;return*this;}
    inline mint&operator*=(const mint&a){v=v*(long long)a.v%mdn;return*this;}
    inline mint&operator/=(const mint&a){v=v*(long long)(~a).v%mdn;return*this;}
    friend istream&operator>>(istream&s,mint&a){s>>a.v;return s;}
    friend ostream&operator<<(ostream&s,const mint&a){s<<a.v;return s;}
};
int n,m,l;
int a[maxn];
mint dp[maxn];
int rep[maxn];
int cnt[maxn];
int l1[maxn],l2[maxn],cl[maxn];
mint sfx[maxn];

namespace nqio{const unsigned R=4e5,W=4e5;char*a,*b,i[R],o[W],*c=o,*d=o+W,h[40],*p=h,y;bool s;struct q{void r(char&x){x=a==b&&(b=(a=i)+fread(i,1,R,stdin),a==b)?-1:*a++;}void f(){fwrite(o,1,c-o,stdout);c=o;}~q(){f();}void w(char x){*c=x;if(++c==d)f();}q&operator>>(char&x){do r(x);while(x<=32);return*this;}q&operator>>(char*x){do r(*x);while(*x<=32);while(*x>32)r(*++x);*x=0;return*this;}template<typename t>q&operator>>(t&x){for(r(y),s=0;!isdigit(y);r(y))s|=y==45;if(s)for(x=0;isdigit(y);r(y))x=x*10-(y^48);else for(x=0;isdigit(y);r(y))x=x*10+(y^48);return*this;}q&operator<<(char x){w(x);return*this;}q&operator<<(char*x){while(*x)w(*x++);return*this;}q&operator<<(const char*x){while(*x)w(*x++);return*this;}template<typename t>q&operator<<(t x){if(!x)w(48);else if(x<0)for(w(45);x;x/=10)*p++=48|-(x%10);else for(;x;x/=10)*p++=48|x%10;while(p!=h)w(*--p);return*this;}}qio;}using nqio::qio;

int main(){
    qio>>n>>m>>l;
    for(int i=1;i<=n;i++)qio>>a[i],a[i]+=a[i]==-1;
    int div=0,lst;
    for(int i=n;i>=1;i--){
        if(a[i])div+=1==++cnt[a[i]],lst=a[i];
        if(a[i+l]&&i+l<=n)div-=0==--cnt[a[i+l]];
        if(div==0)rep[i]=0;
        else if(div==1)rep[i]=lst;
        else rep[i]=-1;
    }
    for(int i=1;i<=l;i++){
        int cc=0;int lp1=i,lp2=i,fp=i;
        l1[i]=l2[i]=i;
        for(int j=i;j<=n;j+=l){
            if(rep[j]==-1){
                lp1=lp2=j+l;cc=0;
            }else if(rep[j]!=0){
                if(cc==0){cc=rep[j];lp2=j+l;}
                else{
                    if(cc!=rep[j]){
                        lp1=fp;lp2=j+l;cc=rep[j];
                    }else lp2=j+l;
                }
                fp=j+l;
            }
            l1[j+l]=lp1;l2[j+l]=lp2;cl[j+l]=cc;
        }
    }
    //for(int i=1;i<=n;i++)cerr<<l1[i]<<' '<<l2[i]<<' '<<cl[i]<<endl;
    dp[n+1]=1;
    for(int i=n;i>=1;i--){
        if(a[i]){
            mint w=dp[i+1];
            if(rep[i]!=-1)w-=dp[i+l];
            if(cl[i]==a[i]){
                sfx[i+l]+=w;
                sfx[l1[i]]-=w;
            }else{
                sfx[i+l]+=w;
                sfx[l2[i]]-=w;
            }
        }else{
            if(rep[i]==0){
                mint w=dp[i+1]-dp[i+l];
                sfx[i+l]+=w*mint(m);
                sfx[l2[i]]-=w*mint(m-1);
                sfx[l1[i]]-=w;
            }else if(rep[i]!=-1){
                mint w1=dp[i+1],w2=dp[i+1]-dp[i+l];
                if(cl[i]==rep[i]){
                    sfx[i+l]+=w1*mint(m-1)+w2;
                    sfx[l2[i]]-=w1*mint(m-1);
                    sfx[l1[i]]-=w2;
                }else{
                    sfx[i+l]+=w1*mint(m-1)+w2;
                    sfx[l2[i]]-=w1*mint(m-2)+w2;
                    sfx[l1[i]]-=w1;
                }
            }else{
                mint w=dp[i+1];
                sfx[i+l]+=w*mint(m);
                sfx[l2[i]]-=w*mint(m-1);
                sfx[l1[i]]-=w;
            }
        }
        dp[i]=sfx[i+l];
        sfx[i]+=sfx[i+l];
        //cerr<<dp[i]<<endl;
        //cerr<<i<<' '<<cl[i]<<' '<<l1[i]<<' '<<l2[i]<<endl;
    }
    qio<<int(dp[1]);
    return 0;
}