HDU - 5769 Substring （后缀数组）

题目链接

题目大意

给出一个字符串s，与一个字符c，求出包含字符c的不同子串数量

思路

只要按照求不同子串数量的思路，特殊处理一下，删掉不包含字符c的子串，就OK啦
我使用的是后缀数组
下面代码带详细注释

代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <set>

using namespace std;
#define mt(x, k) memset((x), (k), sizeof(x))
const int mc = 1e5 + 10;
char s[mc];
int y[mc], x[mc], cnt[mc], sa[mc], rk[mc], height[mc], wt[30];
int n, m;

//因为要输入多组样例，所以记得每次初始化一下，这里是用来初始化的函数
void sa_init() {
    mt(cnt, 0);
    mt(x, 0);
    mt(y, 0);
    mt(sa, 0);
    mt(rk, 0);
    mt(height, 0);
}

//后缀数组模板，不解释了
void get_SA() {
    for (int i = 1; i <= n; ++i) ++cnt[x[i] = s[i]];
    for (int i = 2; i <= m; ++i) cnt[i] += cnt[i - 1];
    for (int i = n; i >= 1; --i) sa[cnt[x[i]]--] = i;
    for (int k = 1; k <= n; k <<= 1) {
        int num = 0;
        for (int i = n - k + 1; i <= n; ++i) y[++num] = i;
        for (int i = 1; i <= n; ++i) if (sa[i] > k) y[++num] = sa[i] - k;
        for (int i = 1; i <= m; ++i) cnt[i] = 0;
        for (int i = 1; i <= n; ++i) ++cnt[x[i]];
        for (int i = 2; i <= m; ++i) cnt[i] += cnt[i - 1]; //第一关键字排名为1~i的数有多少个
        for (int i = n; i >= 1; --i) sa[cnt[x[y[i]]]--] = y[i], y[i] = 0;
        swap(x, y);
        x[sa[1]] = 1;
        num = 1;
        for (int i = 2; i <= n; ++i)
            x[sa[i]] = (y[sa[i]] == y[sa[i - 1]] && y[sa[i] + k] == y[sa[i - 1] + k]) ? num : ++num;
        if (num == n) break;
        m = num;
    }
}

//求高度数组的模板，也不解释
void get_height() {
    int k = 0;
    for (int i = 1; i <= n; ++i) rk[sa[i]] = i;
    for (int i = 1; i <= n; ++i) {
        if (rk[i] == 1) continue;//第一名height为0
        if (k) --k;//h[i]>=h[i-1]-1;
        int j = sa[rk[i] - 1];
        while (j + k <= n && i + k <= n && s[i + k] == s[j + k]) ++k;
        height[rk[i]] = k;//h[i]=height[rk[i]];
    }
}

int main() {
    int t;
    cin >> t;
    for (int k = 1; k <= t; k++) {
        char c[5];
        sa_init();//处理每组样例的时候进行初始化
        scanf("%s %s", c, s + 1);//以字符串读入第一个字符，可以避免处理空格换行之类的问题
        n = strlen(s + 1);
        m = 124;
        get_SA();
        get_height();
        //以上四行都是后缀数组最基本的东西
        //用一个set存储字符c所在的位置
        set<int> st;
        for (int i = 1; i <= n; ++i) {
            if (s[i] == c[0]) st.insert(i);
        }
        //用ans存储结果
        long long ans = 0;
        for (int i = 1; i <= n; i++) {
            //找到当前后缀中，字符c所在的位置
            set<int>::iterator it = st.lower_bound(sa[i]);
            //如果找到了
            if (it != st.end()) {
                //这里是找不同子串的方法
                //首先我们要找有多少个子串
                //比如字符串"baba"，寻找包含字符'a'的不同子串
                //以第一个字符b开头，有4个子串
                //以第二个字符a开头，有3个子串，以此类推
                //可以通过肉眼观察法，得出一个结论，一个后缀有几个字符，就有几个“以当前后缀第一个字符开头”的子串
                //所以用n - sa[i]，因为模板的问题，所以我需要n + 1 - sa[i]
                //但此时仍然有重复的子串，比如"baba" 和 "ba" 都有一个子串"ba"
                //而高度数组中存储的正是最长公共前缀，我们减去height[i]就能得到不同子串的数量
                ans += n + 1 - sa[i] - height[i];
                if (s[sa[i]] != c[0] && (*it - sa[i] - height[i]) > 0)
                    //对于后缀"ba"，它有一个子串"b"，并不包含字符'a'
                    //所以我们要把这些情况处理掉
                    //*it表示的就是当前后缀中，第一个字符'a'的位置
                    //减去前面不包含'a'的前缀，以及重复的，就是我们想要的结果
                    //但是某些情况下，(*it - sa[i] - height[i])的结果可能为负数，所以需要特殊判断一下
                    //比如这组样例
                    /*
					1
					a
					baba
                     */
                    ans -= (*it - sa[i] - height[i]);
            }
        }
        //至此，我们就解决问题了
        //注意输出格式，我一直用cin和cout的，结果被坑了一下QAQ
        printf("Case #%d: %lld\n", k, ans);
    }
    //debug();
    return 0;
}