kmjp's blog

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

Codeforces #305 Div1 E. Mike and Friends

Eの割に解いてる人が多いのでチャレンジしてみた。
http://codeforces.com/contest/547/problem/E

問題

この国では、電話番号はアルファベット小文字で構成されている。
N人各人の電話番号はs[i]とする。

call(i,j)とはi番の人が電話を掛ける際、j番の人に電話がかかる回数である。
これはs[i]の連続部分文字列のうちs[j]と一致する数に一致する。

ここで以下のクエリQ個が与えられる。
各クエリはL,R,Kからなるので、クエリに対し \sum_{i=L}^R call(i,K)を答えよ。

解法

s[i]をスペース区切りで連結した文字列Sを考える。
S中でi番目の人の電話番号が登場する位置をx[i], x[i]+1, .. , y[i]とする。(S[x[i]..y[i]] = s[i])

各クエリに対し、S[x[L]..y[R]]中にs[K]が何回登場するか求めればそれが解になる。
そこで、事前にSuffixArrayを求めておこう。

SuffixArray中でS[x[k]]を先頭とするsuffixのランクをrとする。
lcp(p)~lcp(q-1)≧|s[K]|で、かつlcp(p-1)<|s[K]|かつlcp(q)<|s[K]|となるp,qをRMQの要領で求める。
すなわち、SA中で先頭|s[K]|文字が一致するsuffix群の範囲を求める。
後はランクがこれらp~qの範囲に入るもので、かつsuffixの先頭位置S[y]がx[L]≦y≦Y[R]となるyの数を求める。

p、qは先に各クエリ毎に求めておいて、後はクエリをL,R順に並べ替え、x[L]≦y≦Y[R]となるyの数を数え上げると良い。

struct SuffixArray {
	int N; vector<int> rank,lcp,sa; string S;
	
	SuffixArray(string S) : S(S){
		int i,h=0; vector<int> tmp,tr;
		N=S.size(); rank.resize(N+1); sa.resize(N+1); tmp.resize(N+1);
		FOR(i,N+1) sa[i]=i, rank[i]=i==N?-1:S[i];
		
		for(int k=1; k<=N; k<<=1) {
			auto pred2 = [k,this](int& a,int &b)->bool{ return (((a+k<=N)?rank[a+k]:-1)<((b+k<=N)?rank[b+k]:-1));};
			auto pred = [pred2,k,this](int& a,int &b)->bool{ return (rank[a]!=rank[b])?(rank[a]<rank[b]):pred2(a,b);};
			int x=0;
			if(k!=1) for(i=1;i<N+1;i++) if(rank[sa[i]]!=rank[sa[x]]) sort(sa.begin()+x,sa.begin()+i,pred2), x=i;
			sort(sa.begin()+x,sa.end(),pred);
			FOR(i,N+1) tmp[sa[i]]=(i==0)?0:tmp[sa[i-1]]+pred(sa[i-1],sa[i]);
			swap(rank,tmp);
		}
		lcp.resize(N+1); tr.resize(N+1);
		FOR(i,N+1) tr[sa[i]]=i;
		FOR(i,N) {
			int j=sa[tr[i]-1];
			for(h=max(h-1,0);i+h<N && j+h<N; h++) if(S[j+h]!=S[i+h]) break;
			lcp[tr[i]-1]=h;
		}
	}
};

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V total(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	V add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	V set(int e,V v) { add(e,v-val[e]);}
	int lower_bound(V val) {
		V tv=0; int i,ent=0;
		for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i);
		return ent;
	}
};

int N,Q;
string S[202000];
string sum;
int ind[202000];
int QL[505050],QR[505050],QK[505050];
int A[205050],B[205050],ret[505050];

int rmq[20][1<<21];
BIT<int,19> bt;
vector<int> evs[505050],eve[505050];



void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	sum.reserve(405000);
	FOR(i,N) {
		cin>>S[i];
		ind[i]=sum.size();
		sum+=S[i]+" ";
	}
	ind[N]=sum.size();
	
	SuffixArray sa(sum);
	FOR(i,sum.size()+1) rmq[0][i]=sa.lcp[i];
	FOR(i,19) FOR(x,sum.size()+1) rmq[i+1][x]=min(rmq[i][x],rmq[i][x+(1<<i)]);
	
	FOR(i,N) {
		A[i]=B[i]=x=sa.rank[ind[i]];
		for(j=19;j>=0;j--) {
			if(A[i]-(1<<j)>=0 && rmq[j][A[i]-(1<<j)]>=S[i].size()) A[i]-=(1<<j);
			if(B[i]+(1<<j)<=sum.size() && rmq[j][B[i]]>=S[i].size()) B[i]+=(1<<j);
		}
	}
	
	FOR(i,Q) {
		cin>>QL[i]>>QR[i]>>QK[i];
		QL[i]--;QR[i]--;QK[i]--;
		evs[ind[QL[i]]].push_back(i);
		eve[ind[QR[i]+1]].push_back(i);
	}
	
	FOR(i,sum.size()+1) {
		FORR(r,evs[i]) ret[r] -= bt.total(B[QK[r]]+1)-bt.total(A[QK[r]]);
		FORR(r,eve[i]) ret[r] += bt.total(B[QK[r]]+1)-bt.total(A[QK[r]]);
		bt.add(sa.rank[i]+1,1);
	}
	
	FOR(i,Q) _P("%d\n",ret[i]);
}

まとめ

RMQをSegTree+二分探索で行ったらTLEした。
そのためダブリングベースのRMQをライブラリ化してみた。