POJ - 3376 Finding Palindromes（拓展kmp+trie）

  1 #include<iostream>
  2 #include<algorithm>
  3 #include<string.h>
  4 #define ll long long
  5 using namespace std;
  6 
  7 ///next[i]： T[i]到T[m - 1]与T(模式串)的最长相同前缀长度；
  8 ///extend[i]： S[i]到S[n - 1]（原串）与T的最长相同前缀长度。
  9 
 10 const int maxn=2000100;
 11 int nt[maxn],ex[maxn],k=1,be[maxn],len[maxn];
 12 bool tmp[maxn][2];
 13 char a[maxn],b[maxn];
 14 
 15 struct Node
 16 {
 17     int val;
 18     int color;
 19     int tree[26];
 20 };
 21 Node z[maxn];
 22 int tot, root;
 23 
 24 int newnode()
 25 {
 26     z[tot].val = 0;
 27     z[tot].color = 0;
 28     memset(z[tot].tree, -1, sizeof(z[tot].tree));
 29     tot++;
 30     return tot-1;
 31 }
 32 
 33 //预处理计算next数组
 34 void GETNEXT(char *str,int len)
 35 {
 36     int i=0,j,po;
 37     nt[0]=len;    //初始化nt[0]
 38     while(str[i]==str[i+1]&&i+1<len) i++;    //计算nt[1]
 39     nt[1]=i;
 40     po=1;       //初始化po的位置
 41     for(i=2;i<len;i++){
 42         if(nt[i-po]+i<nt[po]+po) nt[i]=nt[i-po];    //第一种情况，可以直接得到nt[i]的值
 43         else{    //第二种情况，要继续匹配才能得到nt[i]的值
 44             j=nt[po]+po-i;
 45             if(j<0) j=0;    //如果i>po+nt[po],则要从头开始匹配
 46             while(i+j<len&&str[j]==str[j+i]) j++;   //计算nt[i]
 47             nt[i]=j;
 48             po=i;   //更新po的位置
 49         }
 50     }
 51 }
 52 
 53 //计算extend数组
 54 void EXKMP(char *s1,int len,char *s2,int l2,int s,int flag)      ///s1的后缀和s2的前缀匹配
 55 {
 56     int i=0,j,po;
 57     GETNEXT(s2,len);    //计算子串的next数组
 58     while(s1[i]==s2[i]&&i<l2&&i<len) i++;   //计算ex[0]
 59     ex[0]=i;
 60     po=0;   //初始化po的位置
 61     for(i=1;i<len;i++){
 62         if(nt[i-po]+i<ex[po]+po) ex[i]=nt[i-po];      //第一种情况，直接可以得到ex[i]的值
 63         else{    //第二种情况，要继续匹配才能得到ex[i]的值
 64             j=ex[po]+po-i;
 65             if(j<0) j=0;     //如果i>ex[po]+po则要从头开始匹配
 66             while(i+j<len&&j<l2&&s1[j+i]==s2[j]) j++;   //计算ex[i]
 67             ex[i]=j;
 68             po=i;   //更新po的位置
 69         }
 70     }
 71     for(int i=0;i<l2;i++)
 72         if(ex[i]+i==l2) tmp[i+s][flag]=1;
 73 }
 74 
 75 void insert(char *a,int len,int s)
 76 {
 77     int p=root;
 78     for(int i=0;i<len;i++){
 79         int c=a[i]-'a';
 80         z[p].val+=tmp[i+s][0];
 81         if(z[p].tree[c]==-1) {
 82             z[p].tree[c]=newnode();
 83         }
 84         p=z[p].tree[c];
 85     }
 86     z[p].color++;
 87 }
 88 
 89 int query(char *a,int len,int s)
 90 {
 91     int p=root;
 92     ll ans=0;
 93     for(int i=0;i<len;i++){
 94         int c=a[i]-'a';
 95         p=z[p].tree[c];
 96         if(p==-1) break;
 97         if((i<len-1&&tmp[s+i+1][1])||i==len-1) ans+=z[p].color;
 98     }
 99     if(p!=-1) ans+=z[p].val;
100     return ans;
101 }
102 
103 int main()
104 {
105     ios::sync_with_stdio(false);
106     cin.tie(0);
107     cout.tie(0);
108     int t;
109     cin>>t;
110     int l=0;
111     tot=0;
112     root=newnode();
113     for(int i=0;i<t;i++){
114         cin>>len[i]>>a+l;
115         be[i]=l;
116         l+=len[i];
117         for(int j=0;j<len[i];j++){
118             b[j+be[i]]=a[l-1-j];
119         }
120         EXKMP(a+be[i],len[i],b+be[i],len[i],be[i],0);
121         EXKMP(b+be[i],len[i],a+be[i],len[i],be[i],1);
122         insert(a+be[i],len[i],be[i]);
123     }
124     ll ans=0;
125     for(int i=0;i<t;i++){
126         ans+=query(b+be[i],len[i],be[i]);
127     }
128     cout<<ans<<endl;
129     return 0;
130 }