【BZOJ3439】 Kpm的MC密码 （TRIE+主席树）

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<queue>
using namespace std;
#define Maxn 100010
#define Maxd 20

struct node
{
int x,f,p,dfn,rh;
int son[30];
}t[Maxn];

void upd(int x)
{
t[x].p=0;
memset(t[x].son,0,sizeof(t[x].son));
}

char s[Maxn];
int n,nt[Maxn],tot;
int wh[Maxn],k[Maxn];

int mymax(int x,int y) {return x>y?x:y;}

void init()
{
scanf("%d",&n);
upd(0);tot=0;
for(int i=1;i<=n;i++)
{
scanf("%s",s);
int l=strlen(s);
int now=0;
for(int j=l-1;j>=0;j--)
{
int ind=s[j]-'a'+1;
if(!t[now].son[ind])
{
t[now].son[ind]=++tot;
upd(tot);t[tot].f=now;
}
now=t[now].son[ind];
if(j==0)
{
nt[i]=t[now].p;
t[now].p=i;
wh[i]=now;
}
}
}
for(int i=1;i<=n;i++) scanf("%d",&k[i]);
}

int cnt,df[Maxn];
void dfs(int x)
{
t[x].dfn=t[x].rh=++cnt;df[cnt]=x;
for(int i=1;i<=26;i++) if(t[x].son[i])
{
dfs(t[x].son[i]);
t[x].rh=t[t[x].son[i]].rh;
}
}

int rt[Maxn],sum;
struct hp
{
int son[2],cnt;
}a[Maxn*Maxd];

void dfs2(int l,int x,int dep)
{
if(dep==0) return;
if(a[x].son[0]) dfs2(a[l].son[0],a[x].son[0],dep-1);
else a[x].son[0]=a[l].son[0];
if(a[x].son[1]) dfs2(a[l].son[1],a[x].son[1],dep-1);
else a[x].son[1]=a[l].son[1];
}

void build()
{
sum=0;
a[0].son[0]=a[0].son[1]=a[0].cnt=0;
rt[0]=0;
for(int i=1;i<=tot;i++)
{
rt[i]=++sum;
int l,r;
l=rt[i-1],r=sum;
a[rt[i]].cnt=a[rt[i-1]].cnt;
for(int j=t[df[i]].p;j;j=nt[j])
{
int x=j;
a[rt[i]].cnt++;
l=rt[i-1];r=rt[i];
for(int kk=17;kk>=1;kk--)
{
int ind=x/(1<<kk-1);
x%=(1<<kk-1);
l=a[l].son[ind];
if(!a[r].son[ind])
{
a[r].son[ind]=++sum;
a[sum].cnt=a[l].cnt;
a[sum].son[0]=a[sum].son[1]=0;
}
r=a[r].son[ind];
a[r].cnt++;
}
}
dfs2(rt[i-1],rt[i],17);
}
}

int ffind(int l,int r,int kk)
{
l--;
l=rt[l];r=rt[r];
int ans=0;
if(a[r].cnt-a[l].cnt<kk) return -1;
for(int i=17;i>=1;i--)
{
if(a[a[r].son[0]].cnt-a[a[l].son[0]].cnt>=kk)
{
l=a[l].son[0];
r=a[r].son[0];
}
else
{
kk-=a[a[r].son[0]].cnt-a[a[l].son[0]].cnt;
l=a[l].son[1];
r=a[r].son[1];
ans+=(1<<i-1);
}
}
return ans;
}

int main()
{
init();
cnt=-1;
dfs(0);
build();
for(int i=1;i<=n;i++)
{
printf("%d\n",ffind(t[wh[i]].dfn,t[wh[i]].rh,k[i]));
}
return 0;
}