CSA Balanced Strings

Balanced Strings

Description

对于一个仅由$a,b,c$组成的字符串$S$，我们称这个串是合法的当且仅当对于任意一个$S$的连续子串$T$，满足

$$|f(T,x)-f(T,y)|\leq K\ ,\forall\{x,y\}\subset \{'a','b','c'\}$$
其中$f(T,x)$表示$T$中字符$x$的出现次数，$K$是给定的常数。
求有多少长度为$n$的合法串，答案对P取模。

Data Constraints

$n\leq 10^9\ \ \ \ K\leq 5\ \ \ \ 10^8\leq P\leq 10^9+10$

Solution

考虑只有一种限制$(a,b)$的话怎么做。
设$g(i,a,b)$表示$S$串长度为$i$的前缀中$a$的出现次数与$b$的出现次数只差，那么字符串是合法的，当且仅当，

$$max\{g(i,a,b)|i\in [0,n]\}-min\{g(i,a,b)|i\in [0,n]\}\leq K$$
这样的话 $g(i,a,b)$一定在一个长度不超过 $K$的区间内，然后枚举这个区间$K+1$种不同的可能，就可以矩阵乘法求方案数了。

这时又有一个新的问题，如何保证一个合法串只被算一遍。实际上只需要先求出所有长度为$K$的区间的答案，再减去所有长度为$K-1$的区间的答案就可以了。因为我们可以发现，对于任意一个串，如果它的$g(i,a,b)$构成的区间长度为$len$,那么它在第一次中会被算$K+1-len$次，在第二步中会被减去$K-len$次，这样就只会算一次了。

现在设$F(l_1,l_2,l_3)$表示$g(i,a,b),g(i,a,c),g(i,b,c)$三个数构成的值得区间长度限制分别为$l_1,l_2,l_3$的答案，考虑容斥，经过简单的推导（和上面类似），便可以得到答案为（下式用到了对称性）

$$Answer=F(K,K,K)-3F(K-1,K,K)+3F(K-1,K-1,K)-F(K-1,K-1,K-1)$$

现在考虑$F(l_1,l_2,l_3)$如何求解。
$O(K^3)$枚举三个区间，矩阵乘法中的状态只需记录$g(i,a,b),g(i,a,c)$，因为我们有$g(i,b,c)=g(i,a,c)-g(i,a,b)$，我们可以根据第三个区间的限制去掉非法状态。这样矩阵大小是$O(K^2)$，矩阵乘法的复杂度就是$O(K^6)$，加上快速幂和三个区间的$O(K^3)$枚举，复杂度……

然而事实上不同的转移矩阵只有$O(K)$个。
假设我们枚举的三个区间分别是$[-x,l_1-x],[-y,l_2-y],[-z,l_3-z]$，考虑此时那些状态$(u,v)$是合法的(状态$(u,v)$表示$g(i,a,b)=u-x,g(i,a,c)=v-y)$，显然如果合法必有

$$-z\leq (v-y)-(u-x)\leq l_3-z$$

$$y-z-x\leq v-u\leq y-z-x+l_3$$
也就是说只要 $y-z-x$相同，合法状态就是一样的，于是转移矩阵也是一样的，所以只有 $O(K)$个不同的转移矩阵。
时间复杂度 $O(K^7\ log\ n)$，实现特简单，只不过常数嘛……

Code

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>

#define fo(i,j,l) for(int i=j;i<=l;++i)
#define fd(i,j,l) for(int i=j;i>=l;--i)

using namespace std;
typedef long long LL;
const LL K=9,P=K<<2,N=K*K;

LL ls[N][N],jz[N][N],lj[N][N],ke[P][N][N];
int done[N];

int n,k,zd;
LL mo;

inline void fb()
{
    fo(i,1,zd)fo(l,1,zd){
        ls[i][l]=0;
        fo(j,0,zd)ls[i][l]=(ls[i][l]+jz[i][j]*jz[j][l])%mo;
    }
    fo(i,1,zd)fo(l,1,zd)jz[i][l]=ls[i][l];
}

inline void get()
{
    fo(i,1,zd)fo(l,1,zd){
        ls[i][l]=0;
        fo(j,1,zd)ls[i][l]=(ls[i][l]+lj[i][j]*jz[j][l])%mo;
    }
    fo(i,1,zd)fo(l,1,zd)lj[i][l]=ls[i][l];
}

inline int cou(int x,int y,int p)
{return x*p+y+1;}

inline LL FF(int a,int b,int c)
{
    fo(i,0,3*k)done[i]=0;
    LL ans=0; zd=(a+1)*(c+1);
    fo(x,0,a)fo(y,0,c)fo(z,0,b){
        int zs=y-z-x,sx=zs+2*k;
        if(!done[sx]){
            done[sx]=1;
            int gs=(a+1)*(c+1);
            fo(i,1,gs)fo(l,1,gs)lj[i][l]=jz[i][l]=0;
            fo(i,1,gs)lj[i][i]=1;
            int po=0;
            fo(i,0,a)fo(l,0,c)if(l-i>=zs&&l-i<=zs+b){
                ++po;
                if(i!=a&&l!=c)jz[po][cou(i+1,l+1,c+1)]=1;
                if(i&&(l-i+1)<=zs+b)jz[po][cou(i-1,l,c+1)]=1;
                if(l&&(l-1-i)>=zs)jz[po][po-1]=1;
            }else ++po;
            int uy=n;
            for(;uy;uy>>=1,fb())if(uy&1)get();
            fo(i,1,gs)fo(l,1,gs)ke[sx][i][l]=lj[i][l];
        }
        int wz=x*(c+1)+y+1;
        int po=0;
        fo(i,0,a)fo(l,0,c){
            ++po;
            ans=(ans+ke[sx][wz][po])%mo;
        }
    }
    return ans;
}

int main()
{
    cin>>n>>k>>mo;
    LL ans=FF(k,k,k);
    ans=(ans-3*FF(k-1,k,k)%mo+mo)%mo;
    ans=(ans+3*FF(k-1,k,k-1))%mo;
    ans=(ans-FF(k-1,k-1,k-1)+mo)%mo;
    cout<<ans;
}