【模板】后缀数组SA-IS

常见的后缀数组求法为$\Theta(nlogn)$的倍增求法,也有DC3或后缀自动机$\Theta(n)$构造的。我们今天要讲的也是一种$\Theta(n)$构造后缀数组的算法,SA-IS(SuffixArray-InducedSort)

在讲解SA-IS之前,我们需要知道一些概念、定义、以及性质。

后缀数组

对于后缀数组和字符串的定义与概念在此不做过于复杂的讲解,只简要说明以下几点:

  1. 任意可比较大小的元素构成的序列叫做字符串,记为$str$。
  2. $str$从下标$i$开始到结束的子串叫$str$的后缀,记为$suffix(i)$。一个长度为$len$的字符串有$len$个后缀。
  3. 后缀数组$sa[i]$表示将$str$的所有后缀排序后,第$i$小的后缀在原串中的下标。
  4. 名次数组$rank[]$与高度数组$height[]$不是本次讲解重点。

为了方便起见,我们对字符串结尾增加一个空字符(即将末尾的$’\backslash0’$看作字符串的一部分,下文用$#$表示空字符),所以令$len = strlen(str) + 1;$

后缀类型

$L-type$与$S-type$

我们将后缀分为两个类型,$L$型和$S$型。$L$型表示此后缀$suffix(i)$比$suffix(i+1)$大,$S$型表示此后缀$suffix(i)$比$suffix(i+1)$小。特殊的,末尾的空后缀'\0'为$S$型。

举个例子:

对于字符串AGATGAGATACGCGGT,后缀类型是这样的:

下标 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
字符 A G A T G A G A T A C G C G G T #
后缀类型 S L S L L S L S L S S L S S S L S

对于一个字符串,我们可以用如下方法$\Theta(n)$时间内求出它的所有后缀的后缀类型:

type[len - 1] = S;
for(int i = len - 2; i >= 0; --i) {
  if(str[i] > str[i + 1]) type[i] = L;
  if(str[i] < str[i + 1]) type[i] = S;
  if(str[i] == str[i + 1]) type[i] = type[i + 1];
}

对于$str[i]$与$str[i+1]$不一致的地方,正确性显然。对于$str[i]==str[i+1]$的时候,证明如下:

  • 设$suffix(i)$为$aaX$,则$suffix(i+1)$为$aX$。
  • $a=a$,所以比较$aX$与$X$,$aX$与$X$大小关系即为$type[i+1]$。

引理1:对于两个后缀$A$与$B$,若$A[0]=B[0]$且$A$是$L$型,$B$是$S$型,则$A<B$

证明:设$A=abX$,$B=acY$。

若$a\not=b$且$a\not=c$,则$b<a<c$,所以$A<B$。

若$a=b$且$a\not=c$,则比较$aX$与$cY$,$a<c$,所以$A<B$。

若$a\not=b$且$a=c$,则比较$bX$与$aY$,$b<a$,所以$A<B$。

若$a=b$且$a=c$,则比较$aX$与$aY$,转化为前三种情况。

$LMS-type$

$LMS$类型(LeftMostS-type,下简称$*$型)是一种特殊的$S$型,它的要求是它的左边必须是$L$型(所以$#$一定会是$LMS$类型)。

下标 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
字符 A G A T G A G A T A C G C G G T #
后缀类型 S L S L L S L S L S S L S S S L S
LMS * * * * * *

我们可以$\Theta(n)$时间内求出所有$*$型:

bool isLMS(int loc) {
  if(loc <= 0) return false;
  if(type[loc] != S) return false;
  if(type[loc - 1] != L) return false;
  return true;
}

int LMSSize = 0;
for(int i = 0; i < len; ++i) {
  if(!isLMS(type, i)) continue;
  LMS[LMSSize++] = i;
}

一个类型为$*$的字符叫$LMS$字符,两个相邻$LMS$字符所夹的子串(包含这两个字符)叫$LMS$子串,一个$LMS$字符开始的后缀叫$LMS$后缀。特殊的,$#$也是$LMS$子串。

对于$LMS$子串的比较有特殊约定,不应该只比较字符大小,也需要比较类型。只有当每一个字符与其后缀类型都相同时,这两个$LMS$子串才被称为是相同的。

诱导排序(induced sort)

SA-IS的主要思想是利用排序好的$LMS$后缀进行诱导排序,得出后缀数组。关$LMS$后缀的排序我们一会再讲,现在假设我们已经对其排好序了,存在数组$LMS$里:

下标 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
字符 A G A T G A G A T A C G C G G T #
* * * * * *
LMS 16 9 5 7 2 12

诱导排序利用了桶排序的思想,记录每种字符出现的次数,然后放到对应的位置。

显然,后缀数组中首字母相同的后缀在同一段区域。

根据引理1,同种首字母开头的后缀,类型相同的在同一段区域,且$L$型在$S$型前面,如下表所示:

下标 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# A A A A A C C G G G G G G T T T
type S S S S S S S S L L L L S S L L L

诱导排序大致分为以下几步:

  1. 用桶排序的思想记录每种字符出现的次数,并为每种字符在sa中分配相应的空间(如上表),初始化为$-1$。
  2. 逆序扫描$LMS$,把$LMS$后缀逆序依次放入$sa$中对应的$S$型桶末尾。
  3. 正序扫描$sa$,若$type[sa[i]-1]=L$,则将$sa[i]-1$放入对应的$L$型桶中。
  4. 逆序扫描$sa$,若$type[sa[i]-1]=S$,则将$sa[i]-1$放入对应的$S$型桶中。

代码:

/*
str 字符串
len 字符串长度
sigma 字符集大小
type 类型数组
LMS LMS数组
LMSSize LMS数组大小
全局数组sa 后缀数组
*/
template<typename T>
void inducedSort(T str, int len, int sigma, type_t *type, int *LMS, int LMSSize) {
  int *bucket = new int[sigma];
  //bucket 用于统计每个字符出现的次数
  int *buf = new int[sigma];
  //buf 用于分配sa中的内存
  memset(sa, -1, sizeof(sa[0]) * len);
  memset(bucket, 0, sizeof(bucket[0]) * sigma);
  for(int i = 0; i < len; ++i) {
    ++bucket[str[i]];
  }
  buf[0] = bucket[0];
  for(int i = 1; i < sigma; ++i) {
    buf[i] = buf[i - 1] + bucket[i];
  }
  //此时buf[c]表示字符c在sa中对应内存区的最后一位的下标+1

  //把LMS后缀逆序依次放入sa中对应的S型桶末尾
  for(int i = LMSSize - 1; i >= 0; --i) {
    sa[--buf[str[LMS[i]]]] = LMS[i];
  }

  //重置内存区域
  buf[0] = 0;
  for(int i = 1; i < sigma; ++i) {
    buf[i] = buf[i - 1] + bucket[i - 1];
  }
  //此时buf[c]表示字符c在sa中对应内存区的第一位的下标

  //正序扫描sa
  for(int i = 0; i < len; ++i) {
    if(sa[i] <= 0) continue;
    if(type[sa[i] - 1] != L) continue;
    //将sa[i]-1放入对应的L型桶中
    sa[buf[str[sa[i] - 1]]++] = sa[i] - 1;
  }

  //重置内存区域
  buf[0] = bucket[0];
  for(int i = 1; i < sigma; ++i) {
    buf[i] = buf[i - 1] + bucket[i];
  } 
  //此时buf[c]表示字符c在sa中对应内存区的最后一位的下标+1

  //逆序扫描sa
  for(int i = len - 1; i >= 0; --i) {
    if(sa[i] <= 0) continue;
    if(type[sa[i] - 1] != S) continue;
    //将sa[i]-1放入对应的S型桶中
    sa[--buf[str[sa[i] - 1]]] = sa[i] - 1;
  }

  //内存回收
  delete[] bucket;
  delete[] buf;
  return;
}

让我们分析一下这样做的正确性:

首先逆序扫描$LMS$,把$LMS$后缀逆序依次放入$sa$中对应的$S$型桶末尾。这样做可以保证$LMS$后缀在后缀数组中也是有序的

正序扫描$sa$,若$type[sa[i]-1]=L$,则将$sa[i]-1$放入对应的$L$型桶中。

让我们回忆$LMS$的意义:$LeftMostS-type$,也就是说$LMS$的左侧一定是连续的$L$型,并且所有连续的$L$型字符右侧都有一个$LMS$。

然后,由于只查看$L$型,而$L$型的后缀必定比下一位要大,所以放入$sa$中的位置一定在当前位置$i$之后。

由此可知,$L$型后缀不重不漏放入了$sa$数组里。

那么L型后缀是否正确放入了应在的位置呢?我们简单证明一下:

假设两个已知大小关系的后缀$X<Y$,讨论$aX$与$bY$的大小关系:

  1. 当$a\not=b$时,由于$aX$与$bY$各自放入相应的桶里,互不干扰,所以正确排序。
  2. 当$a=b$时,$X<Y$,所以$aX<bY$。由于从左向右扫描$sa$数组,先扫描到$X$,所以先置入$aX$,正确排序。

对于$S$型的证明与$L$型类似。额外说一句在于此时$LMS$的意义已经消失,可以看作普通的$S$型进行排序。

于是,我们得到了该字符串的后缀数组。

但是我们发现还有一个问题刚才被我们跳过去了:LMS后缀的排序。

$LMS$后缀排序

在给所有$LMS$后缀排序之前,我们为所有$LMS$子串按大小离散化,并存入$s1$中。

下标 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
字符 A G A T G A G A T A C G C G G T #
* * * * * *
LMS子串 ATGA AGA ATA ACGC CGGT# #
s1 4 2 3 1 5 0

引理2:$s1$中两个后缀的大小关系,就是$str$中对应的$*$型后缀的大小关系。

证明:显然。 我们可以将$s1$视为是将$*$后缀中不重合的部分进行切割,所以$s1$的后缀就意味着$LMS$后缀,$s1$中两个后缀的大小关系,就是$str$中对应的$*$型后缀的大小关系。

如果向本例这样$s1$中无相同的数值,那么$s1$中后缀大小关系是显然的,可以直接表示出来。

如果$s1$中有相同数值(即存在相同的$LMS$子串),则可递归计算$s1$的后缀大小。

那么我们现在的问题转化成了如何为$LMS$子串离散化(排序):

$LMS$子串的排序

我们依然使用诱导排序,不过往里传的LMS数组并不是排好序的,而是随机的(任意顺序)。

待算法完成,我们获得的$sa$数组中,$LMS$子串之间是排好了序的。

证明:

对于任意后缀$suffix(i)$,本次诱导排序只考虑下标$i$开始到下一个$LMS$字符的部分(称为$LMS$前缀)。显然对于任意$LMS$后缀,有效部分只有第一个字符,所以放入桶中后自然有序。由此,之后的$L$与$S$型排序虽然整体后缀无序,但只考虑$LMS$前缀则有序,最终的$LMS$子串也是有序的。

复杂度分析

至此,SA-IS已经讲完了,让我们分析下它为什么是$\Theta(n)$的时空复杂度:

由于LMS字符个数不会超过字符串长度的一半,所以对于每层SA-IS,时间复杂度为$T(n)$,则有

空间复杂度同理。

总结

  1. 求出$type$数组
  2. 求出$LMS$数组
  3. 排序$LMS$子串
  4. 为$LMS$子串离散化
  5. 排序$LMS$后缀
  6. 诱导出后缀数组

代码

参考资料:
后缀数组及SA-IS算法学习笔记
后缀数组 SA-IS 算法学习心得
A walk through the SA-IS Suffix Array Construction Algorithm