kmjp's blog

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

AtCoder ARC #027 : D - ぴょんぴょんトレーニング

ARC

一見典型DPに見えて地味にメンドイ問題。
http://arc027.contest.atcoder.jp/tasks/arc027_4

問題

1～Nの石が1列に並んでいる。
各石iにはパラメータ1<=H[i]<=10があり、(i+1)～(i+H[i])番の石に飛ぶことができる。

ここでD個のクエリS,Tが与えられる。
S番の石からT番の石に飛ぶ飛び方は何通りあるか。

解法

H[i]の上限をM=10とすると、部分点は以下の手順でO(NMD)で容易に掛ける。
x番の石に到達する飛び方をP[x]とすると、y=1～H[x]に対し P[x+y] += P[x]という配るDPを行うだけで良い。

満点はNが大きいのでO(NMD)ではTLEする。
そこで、どうせ各クエリは似た処理を行うのでこれらを省力化することを考える。

今回は平方分割を行うことにした。
1000の倍数である石番号Aに対し、P[A]～P[A+9]をそれぞれ1と仮置きしたとき、P[A+1000]～P[A+1009]がどうなるかを計算しておく。
実際の処理はP[A]～P[A+9]はもっと大きな値となる場合があるが、それは上記の仮置きした値の線形結合で容易にO(M^2)で計算できる。
この手順でDPを石1000個分ずつ飛ばすことができる。

以下の手順では、まずS以上の最寄りの1000の倍数のAまでは普通に配るDPを行う。
そしてP[A]～P[A+9]を求めたら、あとはTに近い1000の倍数Bまで上記1000飛ばしの処理を行い、またP[B]～P[B+9]の値から配るDPでP[T]を求める。

計算量はおよそO(D*√N*M^2)であるが、Nが大きくMが小さいのでだいぶ実行時間を減らすことができる。

int N;
int H[300101];
int D;
ll dpdp[350][1200][10];
ll mo=1000000007;
ll dp2[305520];

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N;
	FOR(i,N) cin>>H[i];
	FOR(i,300) FOR(x,10) {
		dpdp[i][x][x]=1;
		FOR(y,1000) FOR(j,H[i*1000+y]) dpdp[i][y+j+1][x]=(dpdp[i][y+j+1][x]+dpdp[i][y][x])%mo;
	}
	
	cin>>D;
	while(D--) {
		int S,T;
		cin>>S>>T;
		S--,T--;
		
		for(x=S;x<S+3000;x++) dp2[x]=0;
		dp2[S]=1;
		
		if(T-S>3000) {
			int S2=(S+999)/1000*1000;
			for(x=S;x<S2;x++) FOR(y,H[x]) dp2[x+y+1]=(dp2[x+y+1]+dp2[x])%mo;
			for(S=S2;S+2000<T;S+=1000) {
				FOR(x,10) dp2[S+1000+x]=0;
				FOR(x,10) FOR(y,10) dp2[S+1000+y]+=dp2[S+x]*dpdp[S/1000][1000+y][x]%mo;
				FOR(x,10) dp2[S+1000+x]%=mo;
			}
			for(x=S+10;x<S+4000;x++) dp2[x]=0;
		}
		for(x=S;x<T;x++) FOR(y,H[x]) dp2[x+y+1]=(dp2[x+y+1]+dp2[x])%mo;
		cout << dp2[T] << endl;
	}
}

なお、想定解ではないが実は累積和を使うとO(ND)で書くことができ、これでも7秒台でギリギリ間に合う。

int N;
int H[300101];
int D,S,T;

ll mo=1000000007;
ll ret[300100],sum[300100];

void solve() {
	int f,i,j,k,l,x,y;
	cin>>N;
	FOR(i,N) cin>>H[i];
	cin>>D;
	
	while(D--) {
		cin>>S>>T;
		S--,T--;
		
		FOR(x,11) ret[x+T]=sum[x+T]=0;
		ret[T]=sum[T]=1;
		for(x=T-1;x>=S;x--) {
			ret[x]=sum[x+1]-sum[x+H[x]+1];
			if(ret[x]<0) ret[x]+=mo;
			sum[x]=sum[x+1]+ret[x];
			if(sum[x]>=mo) sum[x]-=mo;
		}
		cout << ret[S] << endl;
	}
}

まとめ

平方分割は思いついたけどうまく時間内に書けなかった。
終わってみたらそこまで長いコードでもないし、ちゃんとやればよかったな。