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*100],ls[N1*100],rs[N1*100],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)>>1; 
    if(x<=mid) update(x,l,mid,ls[rt],w);
    else update(x,mid+1,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)>>1,rt=++tot;
    //sz[rt]=sz[r1]+sz[r2];
    ls[rt]=merge(l,mid,ls[r1],ls[r2]);
    rs[rt]=merge(mid+1,r,rs[r1],rs[r2]);
    pushup(rt);
    return rt;
}
int query(int l,int r,int rt,int K)
{ 
    if(K>sz[rt]) return -1;
    if(l==r) return l;
    int mid=(l+r)>>1;
    if(K>sz[ls[rt]]) return query(mid+1,r,rs[rt],K-sz[ls[rt]]);
    else return query(l,mid,ls[rt],K);
}
}s;

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

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