后缀数组模板 (详细注释)

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[] 里存的是什么。