kmjp's blog

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

TopCoder SRM 595 Div1 Medium LittleElephantAndRGB

最終的には自分で解けたけど、本番中には間に合わなかったな…。
http://community.topcoder.com/stat?c=problem_statement&pm=12113

問題

R,G,BからなるN(<=2500)文字の文字列Sが与えられる。
ここから、0<=a<=b

解法

いろんなやり方がある問題。
a,b,c,dを全て試すとO(N^4)なので当然TLE。

自分はbとcを全通り試す方法をとった。
まず、bとcを固定したとき、S[a..b]がGを連続minGreen個以上含むようなaのとり方の数P[b]と、S[c..d]がGを連続minGreen個以上含むようなdのとり方の数Q[b]を事前にO(N)で計算する。
こうすると、S[a..b]またはS[c..d]のいずれかにGを連続minGreen個以上含む組み合わせは、(S[a..b]にGを連続minGreen個以上含む)*(N-c) + (b)*(S[c..d]にGを連続minGreen個以上含む) - (S[a..b]にGを連続minGreen個以上含む)*(S[c..d]にGを連続minGreen個以上含む)で得られる。

S[b]とS[c]がともにGだった場合、S[a..b]の終端x文字とS[c..d]の頭y文字がGで、かつx=minGreenとなる場合も題意を満たす。
各xとyでループを回してカウントしていっても良いが、そうするとO(N^4)になるので不可。
事前に前処理をして各xおよびyにおける部分和を求めて計算量を落とすとよい。

string S;
int NS[2][2600],L;
ll nb[2600],nc[2600],sum[2600],sum2[2600];

class LittleElephantAndRGB {
	
	public:
	long long getNumber(vector <string> list, int minGreen) {
		ll a,b,c,d,x,y,ret=0;
		int i,j;
		
		S="";
		FOR(i,list.size()) S+=list[i];
		L=S.size();
		ZERO(NS);
		
		FOR(i,L) if(S[i]=='G') NS[0][i]=1+((i==0)?0:NS[0][i-1]);
		for(i=L-1;i>=0;i--) if(S[i]=='G') NS[1][i]=1+((i==L-1)?0:NS[1][i+1]);
		
		ZERO(nb);ZERO(nc);
		x=-1;
		y=L;
		for(b=0;b<L;b++) {
			a=b-(minGreen-1);
			if(a>=0 && NS[1][a]>=minGreen) x=a;
			nb[b]=x;
		}
		for(c=L-1;c>=0;c--) {
			d=c+(minGreen-1);
			if(d<L && NS[0][d]>=minGreen) y=d;
			nc[c]=y;
		}
		
		FOR(b,L) {
			ZERO(sum);
			ZERO(sum2);
			for(i=min(minGreen-1,NS[0][b]);i>=1;i--) sum[i] = sum[i+1] + ((NS[0][b]==i)?(b-nb[b]-i+1):1);
			for(i=1;i<=minGreen-1;i++) sum2[i] = sum2[i-1]+sum[minGreen-i];
			
			for(c=b+1;c<L;c++) {
				// include GGG in either side
				ret += (nb[b]+1)*(L-minGreen-c+1) + (b+1)*(L-nc[c]) - (nb[b]+1)*((L-minGreen)-nc[c]+1); 
				if(S[b]!='G' || S[c]!='G') continue;
				
				if(NS[1][c]<=minGreen) ret += (nc[c]-c-NS[1][c]+1)*sum[minGreen-NS[1][c]] + sum2[NS[1][c]-1];
				else ret += sum2[minGreen-1];
			}
		}
		return ret;
	}
};

まとめ

プランは割とすぐ立ったけど、実装が面倒だった。