square869120Contest #2 : E - 部分文字列

方針は割とすぐ見えたのに、無駄に苦戦。
http://s8pc-2.contest.atcoder.jp/tasks/s8pc_2_e

問題

文字列Sの部分文字列の集合から重複を除いたものに対し、長さの和を求めよ。

解法

重複を除かなければ|S|*(|S|+1)/2である。
ここから重複を除こう。

Sに対しSuffixArrayとLongest Common Prefixを求めよう。
SA上で隣接するSuffix間のLCPがxとすると、両者は先頭x文字までの部分文字列は一致する。
よって重複分x*(x+1)/2を引けばよい。

int L;
string S;

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;
		}
	}
};

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>S;
	L=S.size();
	ll ret=0;
	for(i=1;i<=L;i++) ret += 1LL*i*(L+1-i);
	
	SuffixArray sa(S);
	stack<pair<int,int>> SS;
	
	ll tot=0;
	sa.lcp[L]=L-sa.sa[L];
	for(i=L-1;i>=0;i--) {
		x=sa.lcp[i],
		ret-=1LL*x*(x+1)/2;
	}
	
	cout<<ret<<endl;
}

まとめ

最初不連続な部分文字列でもいいのかと勘違いして見当違いなDPを書いてしまった。