AtCoder ABC #362 (トヨタ自動車プログラミングコンテスト2024#7) : G - Count Substring Query

TLE連発しすぎた。
https://atcoder.jp/contests/abc362/tasks/abc362_g

問題

文字列Sが与えられる。
また、複数のクエリが与えられる。
各クエリは文字列Tで構成される。
Sの部分文字列として、Tと一致する箇所は何通りか答えよ。

解法

SやTに、アルファベット外の文字をくっ付けて連結された文字列を考え、SuffixArrayを求める。
各Tを、SuffixArray上を二分探索し、Sの部分文字列とのLCPが|T|となる範囲を求めよう。

string S;
int N;
int Q;
string T[505050];
int LL[505050];
int id[505050];

template<typename ST=string>
struct SuffixArray {
	int N; vector<int> rank,lcp,sa,rsa; ST S;

	
	SuffixArray(ST S) : S(S){
		int i,h=0; vector<int> tmp;
		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); rsa.resize(N+1);
		FOR(i,N+1) rsa[sa[i]]=i;
		FOR(i,N) {
			int j=sa[rsa[i]-1];
			for(h=max(h-1,0);i+h<N && j+h<N; h++) if(S[j+h]!=S[i+h]) break;
			lcp[rsa[i]-1]=h;
		}
	}
};


template<class V,int NV> class RMQ {
private:
	V table[NV+1][1<<NV];
	int LG[1<<NV];
	int NV2;
public:
	static V const def=-(1<<30);
	V comp(V l,V r){ return min(l,r);};
	RMQ() {
		int i,x;
		NV2=1<<NV;
		LG[1]=0;
		for(i=2;i<NV2;i++) LG[i]=LG[i/2]+1;
		FOR(i,NV) FOR(x,NV2) table[i][x]=def;
	}
	void set(int x,V v){ table[0][x]=v;}
	void build() {
		int i,j,x,y;
		FOR(i,NV) FOR(x,NV2) table[i+1][x]=comp(table[i][x],(x+(1<<i)<NV2)?table[i][x+(1<<i)]:def);
	}
	V query(int L,int R) { //[L,R),
		L=max(0,L), R=min(R,NV2);
		if(R<=L) return def;
		int WL=LG[R-L];
		return comp(table[WL][L],table[WL][R-(1<<WL)]);
	}
	
};

RMQ<int,21> rmq;

int sum[2020202];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>S;
	N=S.size();
	S+="{";
	cin>>Q;
	FOR(i,Q) {
		cin>>T[i];
		LL[i]=T[i].size();
		id[i]=S.size();
		S+=T[i];
		S+="}";
	}
	
	SuffixArray sa(S);
	
	FOR(i,S.size()+1) {
		if(i) sum[i]+=sum[i-1];
		if(sa.sa[i]<N) sum[i]++;
		rmq.set(i,sa.lcp[i]);
	}
	rmq.build();
	
	FOR(i,Q) {
		x=sa.rsa[id[i]];
		int L=x;
		for(j=19;j>=0;j--) if(L-(1<<j)>=0 && rmq.query(L-(1<<j),x)>=LL[i]) L-=1<<j;
		cout<<sum[x]-(L?sum[L-1]:0)<<endl;
	}
	
	
	
}

まとめ

ハッシュ値でやろうとしてTLEしてしまった。