POJ.3729 Facer’s string

1.Problem
3729 -- Facer’s string (poj.org)

2.Idea
Suffix Array
Longest Common Prefix Array

3.Source Code

 const int MAXN = 50010;  
 int n, N, M, k, K;  
 int Rank[MAXN * 2];  
 int tmp[MAXN * 2];  
 int sa[MAXN * 2];  
 int lcp[MAXN * 2];  
 int s[MAXN * 2];  
 bool compare_sa(const int& i, const int& j) {  
      if (Rank[i] != Rank[j])return Rank[i] < Rank[j];  
      else {  
           int ri = i + k <= n ? Rank[i + k] : -1;  
           int rj = j + k <= n ? Rank[j + k] : -1;  
           return ri < rj;  
      }  
 }  
 void construct_sa(const int *S, int *sa) {  
      for (int i = 0; i <= n; i++) {  
           sa[i] = i;  
           Rank[i] = i < n ? S[i] : -1;  
      }  
      for (k = 1; k <= n; k *= 2) {  
           sort(sa, sa + n + 1, compare_sa);  
           tmp[sa[0]] = 0;  
           for (int i = 1; i <= n; i++) {  
                tmp[sa[i]] = tmp[sa[i - 1]] + (compare_sa(sa[i - 1], sa[i]) ? 1 : 0);  
           }  
           for (int i = 0; i <= n; i++) {  
                Rank[i] = tmp[i];  
           }  
      }  
 }  
 void construct_lcp(const int *S, int *sa, int *lcp) {  
      memset(lcp, 0, sizeof(lcp));  
      for (int i = 0; i <= n; i++)Rank[sa[i]] = i;  
      int h = 0;  
      lcp[0] = 0;  
      for (int i = 0; i < n; i++) {  
           int j = sa[Rank[i] - 1];  
           if (h > 0)h--;  
           for (; j + h < n&&i + h < n; h++) {  
                if (S[j + h] != S[i + h])break;  
           }  
           lcp[Rank[i] - 1] = h;  
      }  
 }  
 long long calc(int K) {  
      int A = 0, B = 0;  
      long long res = 0;  
      for (int i = 0; i < n; i++) {  
           if (lcp[i] < K) {  
                if (B > 0)res += A;  
                A = 0; B = 0;  
           }  
           if (sa[i + 1] < N)A++;  
           if (sa[i + 1] > N) B++;  
      }  
      return res;  
 }  
 void solve()  
 {       
      while (scanf("%d%d%d\n", &N, &M, &K) != EOF) {  
           for (int i = 0; i < N; i++) {  
                scanf("%d", &s[i]);  
                s[i]++;  
           }  
           s[N] = '$';  
           for (int i = N + 1; i < N + M + 1; i++) {  
                scanf("%d", &s[i]);  
                s[i]++;  
           }  
           n = N + M + 1;  
           s[n] = 0;  
           construct_sa(s, sa);  
           construct_lcp(s, sa, lcp);  
           printf("%lld\n", calc(K) - calc(K+1));  
      }  
 }