(字典树) 洛谷 P2922 [USACO08DEC]Secret Message G

这题要在字典树上需要纪录两个值，一是a[i]以当前节点结束的字符串有多少，二是b[i]经过当前节点且不在当前节点结束的字符串有多少，查询的时候，如果是以当前查询的字符串为前缀，那值就是查询字符串末尾遍历到的节点i的b[i]值，如果是以n条信息中的字符串为前缀，那就加上查询字符串在字典树上遍历的所有ai值，这两项值加起来就是答案

题面

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define maxn 1050000
int tre[maxn][5],pos;int sum[maxn][3];
int s[10500];
void add(int len)
{
    int c=0;
//printf("%s\n",s);
    for(int i=0;i<len;i++)
    {
        int x=s[i];
        if(tre[c][x]==0)
        {
            tre[c][x]=pos++;
            //printf("tre[%d][%c]=%d\n",c,s[i],tre[c][x]);
        }
        c=tre[c][x];
        sum[c][0]++;
    }
    sum[c][1]++;
    sum[c][0]--;
}
void js(int len)
{
    int c=0,ans0=0,ans1=0;
    for(int i=0;i<len;i++)
    {

        int x=s[i];

        if(tre[c][x]==0)
        {
            c=tre[c][x];
            break;
        }
        c=tre[c][x];
        //printf("tre[%d][%d]=%d s0=%d s1=%d ans1=%d\n",c,x,tre[c][x],sum[c][0],sum[c][1],ans1);
        ans1+=sum[c][1];
    }
    ans0=sum[c][0];
    //printf("ans0=%d ans1=%d\n",ans0,ans1);
    printf("%d\n",ans0+ans1);
    //printf("%d\n",sum[c]);
}
int main()
{
    pos=1;
    int n;int m;
    scanf("%d",&n);scanf("%d",&m);
    for(int i=1;i<=n;i++)
    {
        int len;
        scanf("%d",&len);
        for(int j=0;j<len;j++)
        {
            scanf("%d",&s[j]);

        }add(len);
    }


    for(int i=1;i<=m;i++)
    {
        int len;
        scanf("%d",&len);
        for(int j=0;j<len;j++)
        {
            scanf("%d",&s[j]);
        }
        js(len);
    }
}
/*
10 1000
1 0
1 1
2 1 0
2 0 1
2 1 1
3 0 0 0
3 0 0 1
3 0 1 0
3 1 0 1
1 1
3 1 0 0
*/