hdu6153后缀数组或扩展KMP

前两天刷了几题leetcode，感觉挺简单，于是又想刷刷hduoj了。随便打开没做过的一页，找了一题通过人数最多的，就是这道6153.

①.看完题没想太多，觉得应该是后缀数组(多年没刷题的我字符串这一块对后缀数组记忆最深吧)，因为S1和S2长度都一百万，n^2受不了。nlogn应该行。
②.用后缀数组的话，需要会用后缀数组求子串出现次数。如果是任意子串，还不太好办，但是这里的子串只会是后缀，那就好办了——只需看此后缀在所有后缀中的前缀的数量，也即与本后缀的LCP等于len(s)的后缀数量。因为后缀数组是排过序的，只要往后看，直到height值小于len(s)为止。
③.会了②，就好办了。把S1和S2连起来(中间插入一个比如~号，记为S)，求出S的后缀数组SA；又对S2求出后缀数组SA2。
对于每一个S2的后缀，用在SA中求出的出现次数减去SA2中求出的出现次数，就是此后缀在S1中的出现次数(因为它不会在跨S1和S2时出现)。
④.复杂度分析。构造后缀数组最优算法是O(n)，用倍增算法构造后缀数组是O(nlogn)。后面的计算步骤，因为S2长度n为一百万，要计算n次，每次用②中的办法往后找，最坏情况要O(n)。但是不意味着这里的总复杂度只能O(n^2)。因为后缀数组是排好序的，按后缀数组的逆序处理这n次查找，那么第k次查找可以充分利用第k-1次查找的结果，往后滑动。也就是说这n次查找基本不会重叠，所以最后查找计算部分的总复杂度仍为O(n)。
这个地方我举个例子，以字符串aabaaaab~aabaaaab为例，它有17个后缀，在后缀数组中的顺序为：
aaaab
aaaab~aabaaaab
aaab
aaab~aabaaaab
aab
aabaaaab
aabaaaab~aabaaaab
aab~aabaaaab
ab
abaaaab
abaaaab~aabaaaab
ab~aabaaaab
b
baaaab
baaaab~aabaaaab
b~aabaaaab
~aabaaaab
求出后缀baaaab的出现次数为2以后，再求后缀b的出现次数时，因为知道b是baaaab的子串(height值等于b的长度)，可以直接滑到b~aabaaaab进行判断。
那最后整个算法的复杂度还是卡在构造后缀数组部分，如果用倍增，那整个算法的复杂度为O(nlogn)
⑤.具体实现上，因为我要复用两次后缀数组的代码，所以每次预处理出data数组，再处理出d数组。
d数组就是逆序存的每个后缀在整个串中出现的次数。
代码如下：

/*
 * Author    : ben
 */
#include <cstdio>
#include <cstdlib>
#include <cstring>
typedef long long LL;
const int MAXN = 2010000;
char s[MAXN];
int sa[MAXN], height[MAXN], rank[MAXN], N;
int tmp[MAXN], top[MAXN];
void makesa() {
    int i, j, len, na;
    na = (N < 256 ? 256 : N);
    memset(top, 0, na * sizeof(int));
    for (i = 0; i < N; i++) {
        top[rank[i] = s[i] & 0xff]++;
    }
    for (i = 1; i < na; i++) {
        top[i] += top[i - 1];
    }
    for (i = 0; i < N; i++) {
        sa[--top[rank[i]]] = i;
    }
    for (len = 1; len < N; len <<= 1) {
        for (i = 0; i < N; i++) {
            j = sa[i] - len;
            if (j < 0) {
                j += N;
            }
            tmp[top[rank[j]]++] = j;
        }
        sa[tmp[top[0] = 0]] = j = 0;
        for (i = 1; i < N; i++) {
            if (rank[tmp[i]] != rank[tmp[i - 1]]
                    || rank[tmp[i] + len] != rank[tmp[i - 1] + len]) {
                top[++j] = i;
            }
            sa[tmp[i]] = j;
        }
        memcpy(rank, sa, N * sizeof(int));
        memcpy(sa, tmp, N * sizeof(int));
        if (j >= N - 1) {
            break;
        }
    }
}

void lcp() {
    int i, j, k;
    for (j = rank[height[i = k = 0] = 0]; i < N - 1; i++, k++) {
        while (k >= 0 && s[i] != s[sa[j - 1] + k]) {
            height[j] = (k--), j = rank[sa[j] + 1];
        }
    }
}

char S1[MAXN], S2[MAXN];
int data1[MAXN], data2[MAXN];
int d[MAXN];

void makedata(int *data) {
    data[0] = 1;
    for (int i = N - 2; i > 0; i--) {
        int leni = N - 1 - sa[i];
        int j = N - i - 1;
        if (height[i + 1] < leni) {
            data[j] = 1;
        } else {
            int k = i + data[j - 1] + 1;
            while (k < N && height[k] >= leni) {
                k++;
            }
            data[j] = k - i;
        }
//        cout << data[j] << endl;
    }
}

const LL MOD_NUM = 1000000007LL;
int work(int lens1, int lens2) {
    int ans = 0;
    strcpy(s, S2);
    N = lens2 + 1;
    makesa();
    lcp();
    makedata(data1);
    for (int i = 0; i < lens2; i++) {
        int j = lens2 - i - 1;
        d[i] = data1[lens2 - rank[j]];
//        cout << d[i] << endl;
    }
    strcpy(s, S1);
    s[lens1] = '#';
    s[lens1 + 1] = 0;
    strcat(s, S2);
    N = lens1 + lens2 + 2;
    makesa();
    lcp();
    makedata(data2);
    for (int i = 0; i < lens2; i++) {
        int j = N - i - 2;
        d[i] = data2[N - 1 - rank[j]] - d[i];
        ans = ans % MOD_NUM + ((i + 1LL) * d[i]) % MOD_NUM;
//        cout << d[i] << endl;
    }
    return ans % MOD_NUM;
}

int main() {
    int T;
    scanf("%d", &T);
    while (T--) {
        scanf("%s%s", S1, S2);
        int lens1 = strlen(S1);
        int lens2 = strlen(S2);
        printf("%d\n", work(lens1, lens2));
    }
    return 0;
}

 

然而，提交上去，超时了。看来数据量很强。是卡在后缀数组构造的倍增算法上了。但是我手头没有DC3等更优算法的模板。

转念一想，这题通过的人这么多，不可能需要高阶后缀数组算法的。于是回忆了一下还有别的什么算法。想到了扩展KMP。

用扩展KMP解此题的思路主要是要把S1和S2逆序，逆序以后，题目要求的S2的每个后缀就变成前缀了。根据extend数组的定义我们知道，如果extend[i] = x，则表示S2的前x个字符与S1从i开始的x个字符相同，统计extend数组中有多少个x就知道S2的这个前缀在S1中出现的次数。这种统计是可以在线性时间完成的，而前面生成extend数组的时间也为线性，故最后整体复杂度也为线性O(N)。代码如下

/*
 * Author    : ben
 */
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <ctime>
#include <algorithm>
typedef long long LL;
const int MAXN = 1001000;
char S1[MAXN], S2[MAXN];
int d[MAXN];
const LL MOD_NUM = 1000000007LL;
int next[MAXN], extend[MAXN];
void get_next(const char *str, int len){
    // 计算next[0]和next[1]
    next[0] = len;
    int i = 0;
    while(str[i] == str[i + 1] && i + 1 < len) {
        i++;
    }
    next[1] = i;
    int po = 1; //初始化po的位置
    for(i = 2; i < len; i++) {
        if(next[i - po] + i < next[po] + po) { //第一种情况，可以直接得到next[i]的值
            next[i] = next[i - po];
        } else { //第二种情况，要继续匹配才能得到next[i]的值
            int j = next[po] + po - i;
            if(j < 0) {
                j = 0; //如果i > po + next[po],则要从头开始匹配
            }
            while(i + j < len && str[j] == str[j + i]) { //计算next[i]
                j++;
            }
            next[i] = j;
            po = i; //更新po的位置
        }
    }
}
void extend_KMP(const char *str, int lens, const char *pattern, int lenp) {
    get_next(pattern, lenp); // 先计算模式串的next数组
    // 计算extend[0]
    int i = 0;
    while(str[i] == pattern[i] && i < lenp && i < lens) {
        i++;
    }
    extend[0] = i;
    int po = 0; // 初始化po的位置
    for(i = 1; i < lens; i++) {
        if(next[i - po] + i < extend[po] + po) { //第一种情况，直接可以得到extend[i]的值
            extend[i] = next[i - po];
        } else { // 第二种情况，要继续匹配才能得到extend[i]的值
            int j = extend[po] + po - i;
            if(j < 0) {
                j = 0; //如果i > extend[po] + po则要从头开始匹配
            }
            while(i + j < lens && j < lenp && str[j + i] == pattern[j]) { // 计算extend[i]
                j++;
            }
            extend[i] = j;
            po = i; // 更新po的位置
        }
    }
}

int main() {
    int T;
    scanf("%d", &T);
    while (T--) {
        scanf("%s%s", S1, S2);
        int len1 = strlen(S1);
        int len2 = strlen(S2);
        std::reverse(S1, S1 + len1);
        std::reverse(S2, S2 + len2);
        extend_KMP(S1, len1, S2, len2);
        memset(d, 0, sizeof(d));
        for (int i = 0; i < len1; i++) {
//            printf("%d ", extend[i]);
            d[extend[i]]++;
        }
        LL total = 0LL;
        int ans = 0;
        for (int j = len2; j > 0; j--) {
            total = (total + d[j]) % MOD_NUM;
            ans = (ans + total * j) % MOD_NUM;
        }
        printf("%d\n", ans);
//        putchar('\n');
    }
    return 0;
}

最后，在做此题的过程中我还突然发现输入外挂没用了，加上输入外挂后的执行时间比直接用scanf更长。也许是因为现在的oj用上了最新的编译器吧。