kmjp's blog

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

ゆるふわ競プロオンサイト #2 : K - fuzzy search queries

SAの良い練習。
https://www.hackerrank.com/contests/yfkpo2/challenges/fuzzy-search-queries

問題

文字列Sが与えられる。
以下のクエリに答えよ。

クエリを文字列Qとする。SのQと同じ長さの部分文字列のうち、Qと同じものがあるか。
無い場合、Q以降でQに辞書順で最も近いものを答えよ。

解法

辞書順とか部分列とか出てくるので、Suffix Arrayを使おう。
まず、Sと全クエリのQを連結した文字列を作る。
その際、Sの直後はasciiコード順で'z'以降、各Qの後は'a'以前のもので区切ろう。
こうすると、Suffixが同じ場合、Sに含まれるものが最後に来る。

この上で、range minimum queryを求めるSegTreeを作っておき、Suffix Arrayの結果を見てSに含まれる部分に相当する位置に対しSuffix長を設定しておく。
あとは各クエリに対し、Qの先頭がSuffix Array上でp番目に来たとする。
p番目以降で、上記Suffix長が|Q|以下である最寄りの位置を求めよう。
これはSegTree上で二分探索することで求められる。

struct SuffixArray {
	int N; vector<int> rank,lcp,sa,rsa; string S;
	
	SuffixArray(string 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 SegTree_1 {
public:
	vector<V> val;
	static V const def=0;
	V comp(V l,V r){ return max(l,r);};
	
	SegTree_1(){val=vector<V>(NV*2,0);};
	V getval(int x,int y,int l=0,int r=NV,int k=1) { // x<=i<y
		if(r<=x || y<=l) return def;
		if(x<=l && r<=y) return val[k];
		return comp(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1));
	}
	void update(int entry, V v) {
		entry += NV;
		val[entry]=v;
		while(entry>1) entry>>=1, val[entry]=comp(val[entry*2],val[entry*2+1]);
	}
};

SegTree_1<int,1<<20> st;
string S;
int N;
int Q;
string QS[303030];
int ind[303030];


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>S;
	N=S.size();
	S+="{";
	cin>>Q;
	FOR(i,Q) {
		ind[i]=S.size();
		cin>>QS[i];
		S+=QS[i]+"@";
	}
	
	SuffixArray sa(S);
	FOR(i,S.size()) {
		//cout<<i<<" "<<S.substr(sa.sa[i])<<" "<<N-sa.sa[i]<<endl;
		if(sa.sa[i]<N) st.update(i,N-sa.sa[i]);
	}
	FOR(i,Q) {
		x=sa.rsa[ind[i]];
		y=x+(1<<19);
		if(st.getval(x,y)<QS[i].size()) {
			cout<<"NO"<<endl;
		}
		else {
			for(j=18;j>=0;j--) if(st.getval(x,y-(1<<j))>=QS[i].size()) y-=1<<j;
			string T=S.substr(sa.sa[y-1],QS[i].size());
			if(QS[i]==T) {
				cout<<"YES"<<endl;
			}
			else {
				cout<<T<<endl;
			}
		}
	}
}

まとめ

Suffix Arrayを使うことを思いつけば、あとはスムーズか。