后缀数组（Suffix Array）模板及简析——Part 1：构建SA和rank数组

  1 #include <cstdio>
  2 #include <cstring>
  3 #include <algorithm>
  4 
  5 const int maxN=200005;
  6 
  7 int suffixArray[maxN]; //1-based
  8 int suffixRank[maxN]; //1-based
  9 int height[maxN]; //LCP length of SA[i] and SA[i-1]
 10 
 11 char objStr[maxN];
 12 int objCoded[maxN]; //begin with objCoded[1]
 13 int objLen;
 14 
 15 inline int codeChar(char x)
 16 {
 17     return x+1-‘a‘;
 18 } //range:[1..26]
 19 
 20 void input()
 21 {
 22     scanf("%d",&objLen);
 23     for(int i=0;i<objLen;i+=80) scanf("%s",objStr+i);
 24 
 25     for(int i=1;i<=objLen;i++)
 26         objCoded[i]=codeChar(objStr[i-1]);
 27     objCoded[objLen+1]=-1;
 28 }
 29 
 30 const int alphabetSize=26;
 31 
 32 int charCount[alphabetSize+1];
 33 
 34 int secondKey[maxN]; //use suffixArray[] itself as firstKey
 35 int radixCount[maxN];
 36 int suffixTemp[maxN];
 37 
 38 void buildSuffixArray()
 39 {
 40     memset(charCount,0,sizeof(charCount));
 41 
 42     for(int i=1;i<=objLen;i++)
 43         ++charCount[objCoded[i]];
 44 
 45     for(int i=2;i<=alphabetSize;i++)
 46         charCount[i]+=charCount[i-1];
 47 
 48     for(int i=1;i<=objLen;i++)
 49         suffixRank[i]=charCount[objCoded[i]];
 50 
 51     for(int step=1;step<objLen;step<<=1)
 52     {
 53         for(int i=1;i<=objLen;i++)
 54             secondKey[i]=(i+step<=objLen)?suffixRank[i+step]:0;
 55 
 56         memset(radixCount,0,sizeof(radixCount));
 57 
 58         for(int i=1;i<=objLen;i++)
 59             ++radixCount[secondKey[i]];
 60 
 61         for(int i=1;i<=objLen;i++)
 62             radixCount[i]+=radixCount[i-1];
 63 
 64         for(int i=objLen;i;i--)
 65         {
 66             int pos=radixCount[secondKey[i]]--;
 67             suffixTemp[pos]=i;
 68         }
 69 
 70         memset(radixCount,0,sizeof(radixCount));
 71 
 72         for(int i=1;i<=objLen;i++)
 73             ++radixCount[suffixRank[i]];
 74 
 75         for(int i=1;i<=objLen;i++)
 76             radixCount[i]+=radixCount[i-1];
 77 
 78         for(int i=objLen;i;i--)
 79         {
 80             int& cur=suffixTemp[i];
 81             int pos=radixCount[suffixRank[cur]]--;
 82             suffixArray[pos]=cur;
 83         }
 84 
 85         int last=suffixRank[suffixArray[1]];
 86         suffixRank[suffixArray[1]]=1;
 87         bool finished=true;
 88 
 89         for(int i=2;i<=objLen;i++)
 90         {
 91             int& curPos=suffixArray[i];
 92             int& lastPos=suffixArray[i-1];
 93             int temp=suffixRank[curPos];
 94             if(temp==last && secondKey[curPos]==secondKey[lastPos])
 95             {
 96                 suffixRank[curPos]=suffixRank[lastPos];
 97                 finished=false;
 98             }
 99             else
100                 suffixRank[curPos]=i;
101             last=temp;
102         }
103 
104         if(finished) break;
105     }
106 }
107 
108 void buildHeight()
109 {
110     memset(height,0,sizeof(height));
111 
112     for(int i=1;i<=objLen;i++)
113     {
114         int& curRank=suffixRank[i];
115         int& lastRank=suffixRank[i-1];
116         if(curRank>1)
117         {
118             height[curRank]=height[lastRank]?height[lastRank]-1:0;
119             int& lastSA=suffixArray[curRank-1];
120             int& curSA=suffixArray[curRank]; //curSA=i
121             for(;;height[curRank]++)
122             {
123                 if(objCoded[lastSA+height[curRank]] !=
124                         objCoded[curSA+height[curRank]] )
125                     break;
126             }
127         }
128     }
129 }
130 
131 typedef long long int64;
132 
133 #include <iostream>
134 
135 int main()
136 {
137     input();
138     buildSuffixArray();
139     buildHeight();
140 
141     int64 ans=0;
142 
143     for(int i=1;i<=objLen;i++)
144         ans+=(objLen-suffixArray[i]+1)-height[i];
145 
146     std::cout<<ans<<"\n";
147     return 0;
148 }