kmjp's blog

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

CODE FESTIVAL 2016 Qual B : E - Lexicographical disorder

ライブラリ不足で手間取ったけど、まぁこれはきっちり解ききれて良かった。
http://code-festival-2016-qualb.contest.atcoder.jp/tasks/codefestival_2016_qualB_e

問題

アルファベット小文字だけで構成される、N個の異なる文字列S[i]が与えられる。
これについて以下のクエリQ個に答えよ。

各クエリでは、整数kと26個のアルファベットの大小順が与えられる。
アルファベットの大小順として与えられたものを用いてSを辞書順ソートしたとき、S[k]は何番目に来るか答えよ。

解法

Trie木を用いる。
Trie木の各ノードにおいて、その先にS[i]の終点とマッチするものが何個あるか事前にDFSで求めておこう。
S[k]に応じてTrie木のノードを辿るとき、たとえば(0-indexで)d文字目に相当するノードから次に遷移するときは、S[k][d]より大小順で手前のアルファベットで遷移可能なノードにマッチするものは辞書順で先に来る。
先にDFSで各ノードの先に何個S[i]にマッチするものがあるかはわかっているので、S[k][d]以外のノードはそれ以上探索しなくても良い。
この要領でS[k]より辞書順で手前に来る文字列が何個あるか数え上げる。

各クエリに対してはcをアルファベットの個数(26)としてO(|S[k]|*c)の計算量で処理できる。
ただ、このままだと最悪O(Q*|S[k]|*c)かかりTLEする。
そこでTrie木を圧縮し、行先がどうせ1個しかないノードはスキップできるようにしておこう。
こうするとクエリあたりの探索ノードの数は√(|S[k]|)位になるはず。

int N,Q;
vector<string> S;
int O[26];

const int NUMC=26;
class Trie {
public:
	vector<vector<int> > V;
	vector<int> num,P,I,D;
	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.push_back(vector<int>(NUMC+1));
		num.push_back(0);
		P.push_back(0);
		I.push_back(-1);
		D.push_back(0);
		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;
					num.push_back(0);
					P.push_back(cur);
					I.push_back(-1);
					D.push_back(D[cur]+1);
				}
				cur=V[cur][*c+1];
			}
		}
	}
	
	void compress(int cur,int pre) {
		int i,to=0,tar;
		FOR(i,NUMC) if(V[cur][i+1]) to++,tar=V[cur][i+1];
		
		if(to==1 && I[cur]==-1 && cur) {
			V[P[cur]][pre] = tar;
			P[tar]=P[cur];
			compress(tar,pre);
		}
		else {
			FOR(i,NUMC) if(V[cur][i+1]) compress(V[cur][i+1],i+1);
		}
		
	}
	void dfs(int cur) {
		int i;
		FOR(i,NUMC) if(V[cur][i+1]) dfs(V[cur][i+1]), num[cur] += num[V[cur][i+1]];
		
	}
	
	int query(int tar,int cur) {
		int ret=0,i;
		if(I[cur]==tar) return 1;
		char c=S[tar][D[cur]];
		FOR(i,NUMC) if(V[cur][i+1]) {
			if(O[i]<O[c]) ret += num[V[cur][i+1]];
			if(O[i]==O[c]) ret += query(tar,V[cur][i+1]);
		}
		return ret+(I[cur]>=0);
	}
	
};

Trie T;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>s;
		FORR(c,s) c-='a';
		S.push_back(s);
	}
	
	T.create(S);
	FOR(i,N) T.I[T.find(S[i])]=i, T.num[T.find(S[i])]++;
	T.compress(0,-1);
	T.dfs(0);
	
	cin>>Q;
	while(Q--) {
		cin>>x>>s;
		FOR(i,26) O[s[i]-'a']=i;
		cout<<T.query(x-1,0)<<endl;
	}
}

まとめ

Trie木の圧縮(?)を始めてやったのでバグを作りこんで手間取った。