AtCoder ABC #258 : Ex - Odd Steps

これはどうにか解けた。
https://atcoder.jp/contests/abc258/tasks/abc258_h

問題

整数集合Aと、整数値Sが与えられる。
以下を満たす数列は何通りあるか。

• 各要素は奇数の正整数である。
• 合計はSである。
• prefix sumがAの要素と1つでも一致してはならない。

解法

この問題を以下のように言い換える。

• S段の階段を上ることを考える。この時、1段または2段登ることを繰り返して上る。
• 1段上った段階で、Aに含まれる段数を踏むことは許されない。

後者の条件を除けば、フィボナッチ数列の値を行列累乗で求める典型の問題である。
それを応用しよう。
f(n,k) := 計n段上ったとき、最後にk段上ってきたような上り方

とすると、

• n∉Aの場合
  • f(n,1) = f(n-1,1)+f(n-1,2)
  • f(n,2) = f(n-2,1)+f(n-2,2)
• n∈Aの場合
  • f(n,1) = 0
  • f(n,2) = f(n-2,1)+f(n-2,2)

と状態遷移できる。
この状態遷移を4次正方行列で表現し、S個分掛け合わせよう。
S-|A|個の行列は前者で、|A|の個数だけ後者の行列があるので、前者は行列累乗を使い高速化する。

int N;
ll A[101010],S;
const ll mo=998244353;

const int MAT=4;
struct Mat { ll v[MAT][MAT]; Mat(){ZERO(v);};};

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	ll mo2=4*mo*mo;
	int x,y,z; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(x,n) FOR(z,n) FOR(y,n) {
		r.v[x][y] += a.v[x][z]*b.v[z][y];
		if(r.v[x][y]>mo2) r.v[x][y] -= mo2;
	}
	FOR(x,n) FOR(y,n) r.v[x][y]%=mo;
	return r;
}

Mat powmat(ll p,Mat a,int n=MAT) {
	int i,x,y; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(i,n) r.v[i][i]=1;
	while(p) {
		if(p%2) r=mulmat(r,a,n);
		a=mulmat(a,a,n);
		p>>=1;
	}
	return r;
}



void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>S;
	
	Mat X,Y,E;
	X.v[2][0]=X.v[3][1]=X.v[1][2]=X.v[1][3]=1;
	Y=X;
	X.v[0][0]=1;
	X.v[0][1]=1;
	
	FOR(i,4) E.v[i][i]=1;
	
	ll prev=0;
	FOR(i,N) {
		cin>>A[i];
		
		Mat p=powmat(A[i]-1-prev,X);
		Mat q=mulmat(Y,p);
		E=mulmat(q,E);
		prev=A[i];
	}
	Mat p=powmat(S-prev,X);
	E=mulmat(p,E);
	
	ll ret=E.v[0][0];
	cout<<ret<<endl;
}

まとめ

似たような行列累乗をいくつかの区間で区切るの、前にどこかで見たことある。Codeforcesかな？