BZOJ 3439 Kpm的MC密码 (Trie树+线段树合并)

题面

先把每个串反着插进$Trie$树

每个节点的子树内，可能有一些节点是某些字符串的开头

每个节点挂一棵权值线段树，记录这些节点对应的原来字符串的编号

查询的时候在线段树上二分即可

为了节省空间，使用线段树合并

 #include <vector>
 #include <cstdio>
 #include <cstring>
 #include <algorithm>
 #define N1 100100
 #define M1 (N1*3)
 #define idx(X) (X-'a')
 using namespace std;

 char str[M1];
 int a[N1],n;

 struct SEG{
 int sz[N1*],ls[N1*],rs[N1*],root[M1],tot;
 inline void pushup(int rt){ sz[rt]=sz[ls[rt]]+sz[rs[rt]]; }
 void update(int x,int l,int r,int &rt,int w)
 {
     if(!rt) rt=++tot;
     if(l==r){ sz[rt]+=w; return; }
     int mid=(l+r)>>;
     if(x<=mid) update(x,l,mid,ls[rt],w);
     else update(x,mid+,r,rs[rt],w);
     pushup(rt);
 }
 int merge(int l,int r,int r1,int r2)
 {
     if(!r1||!r2) return r1+r2;
     int mid=(l+r)>>,rt=++tot;
     //sz[rt]=sz[r1]+sz[r2];
     ls[rt]=merge(l,mid,ls[r1],ls[r2]);
     rs[rt]=merge(mid+,r,rs[r1],rs[r2]);
     pushup(rt);
     return rt;
 }
 int query(int l,int r,int rt,int K)
 {
     if(K>sz[rt]) return -;
     if(l==r) return l;
     int mid=(l+r)>>;
     if(K>sz[ls[rt]]) return query(mid+,r,rs[rt],K-sz[ls[rt]]);
     else return query(l,mid,ls[rt],K);
 }
 }s;

 namespace Trie{
 int ch[M1][],dep[M1],fa[M1],tot;
 vector<int>ed[M1];
 void insert(int len,int id,int *pos)
 {
     int x=,c,i;
     for(i=len;i>=;i--)
     {
         c=idx(str[i]);
         if(!ch[x][c])
             ch[x][c]=++tot,dep[tot]=dep[x]+,fa[tot]=x;
         x=ch[x][c];
     }
     ed[x].push_back(id);
     pos[id]=x;
 }
 int hs[M1],que[M1];
 void build()
 {
     int i,j,id,x;
     for(i=;i<=tot;i++) hs[dep[i]]++;
     for(i=;i<=tot;i++) hs[i]+=hs[i-];
     for(i=;i<=tot;i++) que[hs[dep[i]]--]=i;
     for(i=tot;i>=;i--)
     {
         x=que[i];
         for(j=;j<ed[x].size();j++)
         {
             id=ed[x][j];
             s.update(id,,n,s.root[x],);
         }
     }
     for(i=tot;i>=;i--)
     {
         x=que[i];
         s.root[fa[x]]=s.merge(,n,s.root[x],s.root[fa[x]]);
     }
 }
 };

 int pos[M1];
 int main()
 {
     scanf("%d",&n);
     int i,j,k,cnt,L;
     for(i=;i<=n;i++)
     {
         scanf("%s",str+);
         L=strlen(str+);
         Trie::insert(L,i,pos);
     }
     Trie::build();
     for(i=;i<=n;i++)
     {
         scanf("%d",&a[i]);
         printf("%d\n",s.query(,n,s.root[pos[i]],a[i]));
     }
     return ;
 }