HDU 6629 string matching（拓展kmp）

 1 #define bug(x) cout<<#x<<" is "<<x<<endl
 2 #define IO std::ios::sync_with_stdio(0)
 3 #include <bits/stdc++.h>
 4 #define iter ::iterator
 5 using namespace  std;
 6 typedef long long ll;
 7 typedef pair<ll,ll>P;
 8 typedef pair<P,P>P1;
 9 #define mk make_pair
10 #define pb push_back
11 #define se second
12 #define fi first
13 #define rs o<<1|1
14 #define ls o<<1
15 const int N=1e6+5;
16 ll mod=998244353;
17 int T,n;
18 int next1[N],ex[N];
19 
20 void GETNEXT(char *str)
21 {
22     int i=0,j,po,len=strlen(str);
23     next1[0]=len;//初始化next[0]
24     while(str[i]==str[i+1]&&i+1<len)//计算next[1]
25         i++;
26     next1[1]=i;
27     po=1;//初始化po的位置
28     for(i=2; i<len; i++)
29     {
30         if(next1[i-po]+i<next1[po]+po)//第一种情况，可以直接得到next[i]的值
31             next1[i]=next1[i-po];
32         else//第二种情况，要继续匹配才能得到next[i]的值
33         {
34             j=next1[po]+po-i;
35             if(j<0)j=0;//如果i>po+next[po],则要从头开始匹配
36             while(i+j<len&&str[j]==str[j+i])//计算next[i]
37                 j++;
38             next1[i]=j;
39             po=i;//更新po的位置
40         }
41     }
42 }
43 //计算extend数组
44 void EXKMP(char *s1,char *s2)
45 {
46     int i=0,j,po,len=strlen(s1),l2=strlen(s2);
47     GETNEXT(s2);//计算子串的next数组
48     while(s1[i]==s2[i]&&i<l2&&i<len)//计算ex[0]
49         i++;
50     ex[0]=i;
51     po=0;//初始化po的位置
52     for(i=1; i<len; i++)
53     {
54         if(next1[i-po]+i<ex[po]+po)//第一种情况，直接可以得到ex[i]的值
55             ex[i]=next1[i-po];
56         else//第二种情况，要继续匹配才能得到ex[i]的值
57         {
58             j=ex[po]+po-i;
59             if(j<0)j=0;//如果i>ex[po]+po则要从头开始匹配
60             while(i+j<len&&j<l2&&s1[j+i]==s2[j])//计算ex[i]
61                 j++;
62             ex[i]=j;
63             po=i;//更新po的位置
64         }
65     }
66 }
67 
68 char s[N];
69 int main(){
70     scanf("%d",&T);
71     while(T--){
72         scanf("%s",s);
73         //getNext(s);
74         EXKMP(s,s);
75         int len=strlen(s);
76         ll ans=0;
77         for(int i=1;i<len;i++){
78             //bug(extend[i]);
79             ans+=ex[i]+1;
80             if(ex[i]+i>=len)ans--;
81         }
82         printf("%lld\n",ans);
83         for(int i=0;i<=len+2;i++){
84             next1[i]=ex[i]=0;
85         }
86     }
87 }