Codeforces 802I - Fake News(hard) 后缀数组+单调栈

Codeforces 802I - Fake News(hard) 后缀数组+单调栈

题目思路

要求求题目不同子串出现次数的平方和。看到本质不同的子串就可以想到SAM+DP或SA+单调栈。这里记录一下SA+单调栈的解决思路。

首先对字符串$s$建立后缀数组，并求出$height$数组，考虑$height$数组的意义：
$$height[i]=LCP(sa[i],sa[i-1])$$
那么显然当$LCP(sa[i],sa[i-1])<LCP(sa[i-1],sa[i-2])$时，某个公共前缀的贡献终止了，此时应当将该公共前缀的贡献统计进答案。这该如何理解？考虑现在$height[i](LCP)$递增：

如果秃然遇到了一个更小的$LCP$：

那么此时两个阴影覆盖区域的后缀的前缀都终止了贡献。考虑如何统计答案：将答案$n^2$拆分为$2\times (i-1)+1$，然后单独对$2\times (i-1)$的部分进行计数。

维护一个$height$递增的单调栈，每次弹元素的时候统计贡献：贡献等于$2\times \frac{n\times (n+1)}{2}\times len$，统计进答案即可。

Code

#include <bits/stdc++.h>
#define int long long
#define endl '\n'
using namespace std;

const int N = 1e6 + 10;

namespace SA {
    int sa[N], tax[N], pool1[N], pool2[N], height[N];
    int *rnk = pool1, *tp = pool2;
    void sort(int n, int m) {
        for(int i = 0; i <= m; i++) tax[i] = 0;
        for(int i = 1; i <= n; i++) tax[rnk[i]]++;
        for(int i = 1; i <= m; i++) tax[i] += tax[i - 1];
        for(int i = n; i >= 1; i--) sa[tax[rnk[tp[i]]]--] = tp[i];
    }

    void build(string str, int n, int m) {
        for(int i = 1; i <= n; i++) rnk[i] = str[i] - '0' + 1, tp[i] = i;
        sort(n, m);
        for(int w = 1, p = 0; p < n; w <<= 1, m = p) {
            p = 0;
            for(int i = 1; i <= w; i++) tp[++p] = n - w + i;
            for(int i = 1; i <= n; i++) if(sa[i] > w) tp[++p] = sa[i] - w;
            sort(n, m);
            swap(tp, rnk);
            rnk[sa[1]] = p = 1;
            for(int i = 2; i <= n; i++) rnk[sa[i]] = (tp[sa[i - 1]] == tp[sa[i]] && tp[sa[i - 1] + w] == tp[sa[i] + w]) ? p : ++p;
        }
        int k = 0;
        for(int i = 1; i <= n; i++) {
            if(rnk[i] == 1) continue;
            if(k) --k;
            for(int j = sa[rnk[i] - 1]; str[i + k] == str[j + k]; ++k);
            height[rnk[i]] = k;
        } // Get the height array (height[i] = lcp(sa[i], sa[i - 1]))
    } 

    void reset() {
        memset(sa, 0, sizeof(sa));
        memset(rnk, 0, sizeof(pool1));
        memset(tp, 0, sizeof(pool2));
        memset(tax, 0, sizeof(tax));
        rnk = pool1, tp = pool2;
    }
}

using SA::sa;
using SA::rnk;
using SA::height;

inline void solve() {
    string s; cin >> s;
    int n = s.size();
    SA::build('@' + s, n, 233);
    int ans = n * (n + 1) / 2;
    stack<int> stk;
    for(int i = 2; i <= n + 1; i++) {
        while(stk.size() && height[stk.top()] > height[i]) {
            int y = height[stk.top()]; stk.pop();
            int x = height[i];
            int z = (stk.size() ? height[stk.top()] : 0);
            int len = i - (stk.size() ? stk.top() : 1);
            ans += (y - max(x, z)) * (len * len - len);
        }
        stk.emplace(i);
    }
    cout << ans << endl;
    SA::reset();
}

signed main(){
    ios_base::sync_with_stdio(false), cin.tie(0);
    int t = 1; cin >> t;
    while(t--) solve();
    return 0;
}