POJ 3415 Common Substrings

  1 #include <iostream>
  2 #include <cstdio>
  3 #include <cstring>
  4 using namespace std;
  5 const int maxn = 100010;
  6 typedef long long LL;
  7 int ha(char ch){
  8     if(ch >= ‘a‘ && ch <= ‘z‘) return ch - ‘a‘;
  9     return ch - ‘A‘ + 26;
 10 }
 11 struct SAM{
 12     struct node{
 13         int son[52],f,len;
 14         void init(int _len){
 15             f = -1;
 16             len = _len;
 17             memset(son,-1,sizeof son);
 18         }
 19     }e[maxn<<1];
 20     int cnt[maxn<<1],lazy[maxn<<1],tot,last;
 21     int c[maxn<<1],sa[maxn<<1];
 22     int newnode(int len = 0){
 23         e[tot].init(len);
 24         return tot++;
 25     }
 26     void init(){
 27         tot = last = 0;
 28         memset(cnt,0,sizeof cnt);
 29         memset(lazy,0,sizeof lazy);
 30         memset(c,0,sizeof c);
 31         newnode();
 32     }
 33     void extend(int c){
 34         int p = last,np = newnode(e[p].len + 1);
 35         while(p != -1 && e[p].son[c] == -1){
 36             e[p].son[c] = np;
 37             p = e[p].f;
 38         }
 39         if(p == -1) e[np].f = 0;
 40         else{
 41             int q = e[p].son[c];
 42             if(e[p].len + 1 == e[q].len) e[np].f = q;
 43             else{
 44                 int nq = newnode();
 45                 e[nq] = e[q];
 46                 e[nq].len = e[p].len + 1;
 47                 e[np].f = e[q].f = nq;
 48                 while(p != -1 && e[p].son[c] == q){
 49                     e[p].son[c] = nq;
 50                     p = e[p].f;
 51                 }
 52             }
 53         }
 54         cnt[last = np] = 1;
 55     }
 56     void update(){
 57         for(int i = tot-1; i > 0; --i)
 58             cnt[e[sa[i]].f] += cnt[sa[i]];
 59     }
 60     void sort(int len){
 61         for(int i = 0; i < tot; ++i) c[e[i].len]++;
 62         for(int i = 1; i <= len; ++i) c[i] += c[i-1];
 63         for(int i = tot-1; i >= 0; --i) sa[--c[e[i].len]] = i;
 64     }
 65     LL solve(char str[],int k,LL ret = 0){
 66         int p = 0,tmp = 0;
 67         for(int i = 0; str[i]; ++i){
 68             int c = ha(str[i]);
 69             if(~e[p].son[c]){
 70                 p = e[p].son[c];
 71                 ++tmp;
 72             }else{
 73                 while(p != -1 && e[p].son[c] == -1) p = e[p].f;
 74                 if(p == -1) {
 75                     p = 0;
 76                     tmp = 0;
 77                 }else{
 78                     tmp = e[p].len + 1;
 79                     p = e[p].son[c];
 80                 }
 81             }
 82             if(tmp >= k){
 83                 ret += (tmp - max(k,e[e[p].f].len + 1) + 1)*cnt[p];
 84                 if(k <= e[e[p].f].len) ++lazy[e[p].f];
 85             }
 86         }
 87         for(int i = tot-1; i > 0; --i){
 88             int u = sa[i];
 89             ret += (LL)lazy[u]*(e[u].len - max(k,e[e[u].f].len + 1) + 1)*cnt[u];
 90             if(k <= e[e[u].f].len) lazy[e[u].f] += lazy[u];
 91         }
 92         return ret;
 93     }
 94 }sam;
 95 char str[maxn];
 96 int main(){
 97     int k;
 98     while(scanf("%d",&k),k){
 99         scanf("%s",str);
100         sam.init();
101         for(int i = 0; str[i]; ++i)
102             sam.extend(ha(str[i]));
103         sam.sort(strlen(str));
104         sam.update();
105         scanf("%s",str);
106         printf("%lld\n",sam.solve(str,k));
107     }
108     return 0;
109 }