kmjp's blog

競技プログラミング参加記です

AtCoder ABC #213 : H - Stroll

似たようなテク、過去に使ったことあるけど忘れてたな…。
https://atcoder.jp/contests/abc213/tasks/abc213_h

問題

N頂点のグラフがある。
あるM通りの2頂点(a,b)の間には無向辺が多数ある。
距離t毎に、その距離と一致する長さの辺が何本あるかが与えられる。

頂点1から移動を開始し、距離Tだけ移動して頂点1に戻る経路は何通りあるか求めよ。
なお、辺の途中で折り返すことはできないとする。

f(v,t) := 距離tだけ移動して頂点vに至る経路の数
とすると、f(1,T)を求めればよい。

また、(a,b)間に距離tの辺がg(a,b,t)本ある場合、
f(b,s+t) += f(a,s) * g(a,b,t)
と平易な遷移となるが、あいにく辺の組み合わせはO(MT)通りあるため、この処理は愚直に行うとO(NMT^2)かかってしまう。

上記式は畳み込みの形式に近いので、畳み込みを使うことを考える。
区間[L,R)において、f(*,0)~f(*,L-1)からf(*,L)~f(*,R-1)に対する寄与分がすでに分かっているとき、区間[L,R)内の移動における寄与分を計算する処理を考える。
L,Rの中点をMとしたとき、以下のようにする。

  • f(*,L)~f(*,M-1)を再帰的に求める。
  • f(*,L)~f(*,M-1) から f(*,M)~f(*,R-1)への寄与分を求める。これは畳み込みで行える。
  • f(*,M)~f(*,R-1)を再帰的に求める。

f(*,0)は確定しているので、区間[0,T+1)を対象に上記処理を行えばよい。

const int mo=998244353;
ll modpow(ll a, ll n = mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}


template<class T> vector<T> fft(vector<T> v, bool rev=false) {
	int n=v.size(),i,j,m;
	
	for(i=0,j=1;j<n-1;j++) {
		for(int k=n>>1;k>(i^=k);k>>=1);
		if(i>j) swap(v[i],v[j]);
	}
	for(int m=2; m<=n; m*=2) {
		T wn=modpow(5,(mo-1)/m);
		if(rev) wn=modpow(wn);
		for(i=0;i<n;i+=m) {
			T w=1;
			for(int j1=i,j2=i+m/2;j2<i+m;j1++,j2++) {
				T t1=v[j1],t2=(ll)w*v[j2]%mo;
				v[j1]=t1+t2;
				v[j2]=t1+mo-t2;
				while(v[j1]>=mo) v[j1]-=mo;
				while(v[j2]>=mo) v[j2]-=mo;
				w=(ll)w*wn%mo;
			}
		}
	}
	if(rev) {
		ll rv = modpow(n);
		FOR(i,n) v[i]=(ll)v[i]*rv%mo;
	}
	return v;
}

template<class T> vector<T> MultPoly(vector<T> P,vector<T> Q,bool resize=false) {
	if(resize) {
		int maxind=0,pi=0,qi=0,i;
		int s=2;
		FOR(i,P.size()) if(norm(P[i])) pi=i;
		FOR(i,Q.size()) if(norm(Q[i])) qi=i;
		maxind=pi+qi+1;
		while(s*2<maxind) s*=2;
		P.resize(s*2);Q.resize(s*2);
		if(s<=16) { //fastpath
			vector<T> R(s*2);
			for(int x=0;x<2*s;x++) for(int y=0;x+y<2*s;y++) (R[x+y]+=P[x]*Q[y])%=mo;
			return R;
		}
	}
	P=fft(P), Q=fft(Q);
	for(int i=0;i<P.size();i++) P[i]=(ll)P[i]*Q[i]%mo;
	return fft(P,true);
}

int N,M,T;
ll dp[10][50000];
int A[10],B[10];
ll P[10][50000];

void dfs(int L,int R) {
	if(L+1>=R) return;
	int m=(L+R)/2;
	dfs(L,m);
	int i,j,x;
	FOR(i,M) {
		vector<ll> Va,Vb,V2;
		for(j=L;j<m;j++) Va.push_back(dp[A[i]][j]);
		for(j=L;j<m;j++) Vb.push_back(dp[B[i]][j]);
		for(j=0;j<R-L;j++) V2.push_back(P[i][j]);
		Va=MultPoly(Va,V2,1);
		Vb=MultPoly(Vb,V2,1);
		for(j=m;j<R;j++) if(j-L<Va.size()) (dp[B[i]][j]+=Va[j-L])%=mo;
		for(j=m;j<R;j++) if(j-L<Vb.size()) (dp[A[i]][j]+=Vb[j-L])%=mo;
	}
	dfs(m,R);
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>T;
	FOR(i,M) {
		cin>>A[i]>>B[i];
		A[i]--;
		B[i]--;
		FOR(j,T) cin>>P[i][j+1];
	}
	dp[0][0]=1;
	dfs(0,T+1);
	
	cout<<dp[0][T]<<endl;
	
}

まとめ

8問制になって典型テクの敷居が上がったな。