kmjp's blog

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

Codeforces ECR #070 : E. You Are Given Some Strings...

これは思いついてしまえば楽。
https://codeforces.com/contest/1202/problem/E

問題

文字列t,sに対し、f(t,s)とはtの部分文字列としてsが現れる回数とする。
文字列Tと、多数の文字列S[i]が与えられる。
(i,j)の全対に対し、f(T,S[i]+S[j])の総和を求めよ。

解法

Tのx文字目までのsuffixがS[i]と一致し、Tのx+1文字目以降のprefixがS[j]と一致するなら解に1が加算される。
さらに言えば、Tのx文字目までのsuffixがS[i]と一致するようなiがp個あり、Tのx+1文字目以降のprefixがS[j]と一致するようなjがq個あるなら、p*qが解に加算される。

よって、Tの各位置までの部分文字列のSuffixが何個S[i]と一致し、各位置以降の部分文字列のPrefixが何個S[j]と一致するかを求めればよい。
これ普通にTrieを作ってAho-Corasickを行いつつ、TもSも逆向きにしたもので同様にAho-Corasickを行えばよい。

const int NUMC=26;
class Trie {
public:
	vector<vector<int> > V;
	int find(string s) {
		int cur=0;
		ITR(it,s) if((cur=V[cur][*it+1])==0) return -1;
		return cur;
	}
	void create(vector<string> S) { // 0 is for backtrack
		V.clear();
		V.push_back(vector<int>(NUMC+1));
		sort(S.begin(),S.end());
		ITR(it,S) {
			int cur=0;
			ITR(c,(*it)) {
				if(V[cur][*c+1]==0) V.push_back(vector<int>(NUMC+1)),V[cur][*c+1]=V.size()-1;
				cur=V[cur][*c+1];
			}
		}
	}
};

class ACmatch_num {
public:
	Trie t;
	vector<int> acc;
	int ma;
	void create(vector<string> S) {
		int i;
		ma=S.size();
		t.create(S);
		acc.clear();
		acc.resize(t.V.size());
		FOR(i,S.size()) acc[t.find(S[i])]++;
		queue<int> Q;
		FOR(i,NUMC) if(t.V[0][i+1]) t.V[t.V[0][i+1]][0]=0, Q.push(t.V[0][i+1]);
		
		while(!Q.empty()) {
			int k=Q.front(); Q.pop();
			FOR(i,NUMC) if(t.V[k][i+1]) {
				Q.push(t.V[k][i+1]);
				int pre=t.V[k][0];
				while(pre && t.V[pre][i+1]==0) pre=t.V[pre][0];
				t.V[t.V[k][i+1]][0]=t.V[pre][i+1];
				acc[t.V[k][i+1]] += acc[t.V[pre][i+1]];
			}
		}
	}
	vector<int> match(string S) {
		int cur=0;
		vector<int> V;
		ITR(it,S) {
			while(cur && t.V[cur][*it+1]==0) cur=t.V[cur][0];
			cur=t.V[cur][*it+1];
			V.push_back(acc[cur]);
		}
		return V;
	}
};

string T,RT;
vector<string> S,RS;
int L,N;
ACmatch_num A,B;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	FORR(c,T) c-='a';
	L=T.size();
	RT=T;
	reverse(ALL(RT));
	
	cin>>N;
	FOR(i,N) {
		cin>>s;
		FORR(c,s) c-='a';
		S.push_back(s);
		reverse(ALL(s));
		RS.push_back(s);
	}
	
	A.create(S);
	B.create(RS);
	auto a=A.match(T);
	auto b=B.match(RT);
	reverse(ALL(b));
	ll ret=0;
	FOR(i,a.size()-1) ret+=1LL*a[i]*b[i+1];
	cout<<ret<<endl;
	
	
}

まとめ

こういうのがスパっと決まると気持ちよいね。