后綴數組模板 (詳細注釋)

2019/12/12 更新：
把代碼整理、優化了一下。

// 要求T具有<運算符
template <typename T>
void sort_suffix(const T *s, int n, vector<int> &suf, vector<int> &h) {
  assert(SZ(suf) >= n && SZ(h) >= n);
  vector<int> rank(n);
  iota(all(suf), 0);
  sort(all(suf), [s](int x, int y) { return s[x] < s[y]; });
  rank[suf[0]] = 0;
  for (int i = 1; i < n; ++i) {
    rank[suf[i]] = rank[suf[i - 1]] + (s[suf[i - 1]] < s[suf[i]]);
  }
  int m = rank[suf[n - 1]] + 1;
  if (m < n) {
    //判斷兩個長為2*len的子串是否相等。
    //先比較第一關鍵字：前一半的rank，再比較第二關鍵字：后一半的rank。
    auto same = [n](const int *rank, int idx1, int idx2, int len) {
      assert(idx1 != idx2);
      return idx1 + len < n && idx2 + len < n && rank[idx1] == rank[idx2] &&
             rank[idx1 + len] == rank[idx2 + len];
    };
    vector<int> cnt(n), suf2(n);
    int len = 1;
    do {
      // 第一步：將子串按第二關鍵字排序，結果存到數組suf2。
      // 這一步利用了上一輪算出的suf[]數組。
      // 這個過程可形象地理解為將子串按順序從數組suf中拿出來放到數組suf2中。
      int p = 0;
      // 先把后一半是空串的子串放在開頭，
      for (int i = n - len; i < n; i++)
        suf2[p++] = i;
      // 接着把后一半不是空串的子串按順序逐個加到末尾。
      for (int i = 0; i < n; i++)
        if (suf[i] >= len)
          suf2[p++] = suf[i] - len;
      // 第二步：將n個長為2*len的子串排序。
      // 計數排序。
      // 每個子串的第一關鍵字是它的前一半的rank。
      for (int i = 0; i < m; i++)
        cnt[i] = 0;
      for (int i = 0; i < n; i++)
        cnt[rank[i]]++;
      for (int i = 1; i < m; i++)
        cnt[i] += cnt[i - 1];
      for (int i = n - 1; i >= 0; i--)
        suf[--cnt[rank[suf2[i]]]] = suf2[i];
      // 第一步和第二步合起來構成一次基數排序。
      //計算rank。
      swap(rank, suf2);
      rank[suf[0]] = 0;
      for (int i = 1; i < n; i++) {
        //此時suf2存的是舊的rank。
        rank[suf[i]] =
            rank[suf[i - 1]] + !same(suf2.data(), suf[i - 1], suf[i], len);
      }
      m = rank[suf[n - 1]] + 1;
      len *= 2;
    } while (m < n);
  }
  //計算高度數組h。
  //對於 i >= 1, h[rank[i]] >= h[rank[i-1]] - 1
  for (int i = 0, lcp = 0; i < n; i++) {
    //循環不變量：每輪循環開始之前，有 h[rank[i]] >= lcp 且 lcp >= 0。
    if (rank[i] > 0) {
      int x = suf[rank[i] - 1] + lcp, y = i + lcp;
      while (x < n && y < n && s[x] == s[y]) {
        ++x, ++y, ++lcp;
      }
      h[rank[i]] = lcp;
      if (lcp > 0)
        --lcp;
    } else {
      // 此時有 h[rank[i]] >= lcp 且 lcp >= 0
      h[rank[i]] = lcp;
    }
  }
}

---

2015年5月初次學習后綴數組，當時是 error202 講的。代碼一直不熟練，導致運用得不熟練，2016年 ACM-ICPC China Finals 就有一道后綴數組的裸題，我們隊並沒有過。現在回想起來，真是荒廢了大把時間，悔之晚矣。

Implementation

后綴數組原理比較好理解, 實現起來也有比較簡單的寫法。下面的實現參考了 2009 年國家集訓隊隊員羅穗騫的論文《后綴數組--處理字符串的有力工具》中的實現。論文中給出的實現比較巧妙，但比較難理解，需要讀者懂得計數排序（counting sort）和基數排序（radix sort）作為預備。

//比較兩個串是否相同
//分別比較兩個關鍵字
bool same(int *rank, int idx1, int idx2, int len){
    return rank[idx1]==rank[idx2] && rank[idx1+len]==rank[idx2+len];
}

// 輸入字符串的末尾要補一個 '0' (C-style string 默認如此), n是字符串的實際長度+1. 
void SA(int *s, int *sa, int *sa2, int *rk, int *cnt, int *hgt, int n, int m){

　　//counting sort
    for(int i=0; i<m; i++) cnt[i]=0;
    for(int i=0; i<n; i++) cnt[rk[i]=s[i]]++;
    for(int i=1; i<m; i++) cnt[i]+=cnt[i-1];
    for(int i=n-1; i>=0; i--) sa[--cnt[rk[i]]]=i;　　//stable sort
  
    for(int len=1; len<n; len*=2){
        //step-1: 計算（填充） sa2[]
        //sa2[]: 按第二關鍵字對后綴排序的結果
        //對后綴按第二關鍵字排序這一步可直接用上一次的sa[]數組, 這個過程可形象地理解為將后綴按順序從一個數組中拿出來放到另一個數組中
        int p=0;
        for(int i=n-len; i<n; i++) sa2[p++]=i;   //第二關鍵字為空的后綴放在最開頭
        //接着放第二關鍵字非空的
        //形象化地理解: 將后綴從一個數組里按順序拿出來放到另一個數組里
        for(int i=0; i<n; i++) 
            if(sa[i]>=len) 
                sa2[p++]=sa[i]-len;

        //step-2 fill sa[], countig sort
        //計數排序: 每個后綴的第一關鍵字 (first key) 為上一次針對長度減半的 (部分) 后綴求出來的rank
        for(int i=0; i<m; i++) cnt[i]=0;
        for(int i=0; i<n; i++) cnt[rk[i]]++;
        for(int i=1; i<m; i++) cnt[i]+=cnt[i-1];
        for(int i=n-1; i>=0; i--) 
            sa[--cnt[rk[sa2[i]]]]=sa2[i];
        //step-1 and step-2 together constitute a radix sort

        //fill rk[]
        swap(rk, sa2);
        rk[sa[0]]=0;
        for(int i=1; i<n; i++)
            //這里, sa2指向的是舊rank數組
            rk[sa[i]]=rk[sa[i-1]]+!same(sa2, sa[i-1], sa[i], len);

        m=rk[sa[n-1]]+1;
        if(m==n) break;
    }

    //CALCULATE hgt[]

    for(int i=0, j, lcp=0; i<n-1; i++){
        lcp?--lcp:0;
        // rk[i]>0
        j=sa[rk[i]-1];    
        for(; s[j+lcp]==s[i+lcp]; lcp++);
        hgt[rk[i]]=lcp;
    }  
}

實現要點

接口

void SA(int *s, int *sa, int *sa2, int *rk, int *cnt, int *hgt, int n, int m)

s: 輸入字符串, 末尾補一個特殊字符 \0
sa: 后綴數組
sa2: 以第二關鍵字對后綴排序所得數組, 輔助數組
cnt: 用於計數排序的輔助數組
rk: rank數組, 輔助數組
n: 輸入字符串 s 的長度, 為其實際長度+1
m: 單字符rank的范圍, 輔助變量.

在輸入字符串 s 末尾補 0 的好處

羅穗騫論文中概括為

為了函數操作的方便

具體來說，有兩處：
一是 same 函數：

bool same(int *rank, int idx1, int idx2, int len){
    return rank[idx1]==rank[idx2] && rank[idx1+len]==rank[idx2+len];
}

由於在字符串的末尾補了一個 \0，rank[idx1] == rank[idx2] 保證了 idx1 + len，idx2 + len 這兩個下標不會越界。

一是在求 hgt 數組時,

    for(int i=0, j, lcp=0; i<n-1; i++){
        lcp?--lcp:0;
        // rk[i]>0
        j=sa[rk[i]-1];    
        for(; s[j+lcp]==s[i+lcp]; lcp++);
        hgt[rk[i]]=lcp;
    }

由於輸入數組末尾補零的緣故, 有 $\mathrm{rk}[n-1]=0$，$\mathrm{rk}[i]>0, (0\le i < n-1)$。
從而 j=sa[rk[i]-1]; 不會越界。

height 數組的求法

height[i]：sa[i] 和 sa[i-1] 兩后綴的最長公共前綴（LCP）的長度。
height[i] 具有如下性質：
$$\mathrm{height}[\mathrm{rank}[i]] \ge \mathrm{height}[\mathrm{rank}[i-1]]-1$$
依據這個性質, 我們可以按 $\mathrm{height}[\mathrm{rank}[0]], \mathrm{height}[\mathrm{rank}[1]], \dots, \mathrm{height}[\mathrm{rank}[n-1]]$ 的順序計算 $\mathrm{height}[]$ 數組, 復雜度為 $O(n)$ .

注意：上面的實現是在 SA() 內部計算的 hgt[] 的。計算 hgt[] 需借助 rank[]，若要在 SA() 外面計算 hgt[]，則須根據求出的 sa[] 再算一遍 rk[]。rk 是作為指針傳入 SA() 的，然而 SA() 中有 swap(rk, sa2); 這一條語句，因此在 SA() 外部並不能斷定 rk[] 里存的是什么。