[SAM模板题] P3975 [TJOI2015] 弦论

题目思路

对原串建立$SAM$，然后分情况考虑：

• $t=0$，代表本质不同即为不同，那么我们需要对每个$SAM$节点$u$维护$u$能够到达的点的数量，即为经过$u$的本质不同的子串的数量。
• $t=1$，代表位置不同即为不同，那么我们考虑对于每个节点$u$计算经过该点位置不同的子串数量。那么我们可以想到这需要先求串$s$的出现次数，也就是$endpos(s)$的大小，而根到$u$表示的串的$endpos$大小为$parenttree$上$u$的子树中的前缀点的数量

不难发现，我们可以维护一个$dp[u]$表示不同限制下的子串数目，以及每个点的点权。对于$t=0$，所有点的点权都是$1$，否则先$dfs$求子树大小作为点权，然后$dfs$求答案即可。

#include <bits/stdc++.h>
#pragma gcc optimize("O2")
#pragma g++ optimize("O2")
#define int long long
#define endl '\n'
using namespace std;

const int N = 1e6 + 10, MOD = 1e9 + 7;

int dp[N], siz[N];

struct SAM{
    
    int ch[N << 1][26], fa[N << 1], len[N << 1], vis[N << 1];
    int last, tot;
    SAM(): last(1), tot(1) {
    }
    inline void extend(int x){
     //*单字符扩展
        int p = last, np = last = ++tot;
        len[np] = len[p] + 1, vis[np] = dp[np] = siz[np] = 1;
        for(; p && !ch[p][x]; p = fa[p]) ch[p][x] = np;
        if(!p) fa[np] = 1;
        else{
    
            int q = ch[p][x];
            if(len[q] == len[p] + 1) fa[np] = q;
            else {
    
                int nq = ++tot;
                memcpy(ch[nq], ch[q], sizeof(ch[nq]));
                fa[nq] = fa[q], fa[np] = fa[q] = nq, len[nq] = len[p] + 1;
                for(; ch[p][x] == q; p = fa[p]) ch[p][x] = nq;
            }
        }
    }
}sam;

vector<int> g[N];

void dfs1(int u){
    
    for(auto v : g[u]){
    
        dfs1(v);
        siz[u] += siz[v];
    }
    dp[u] = siz[u];
}

bool vis[N];

int dfs2(int u){
    
    if(vis[u]) return dp[u];
    vis[u] = true;
    for(int i = 0; i < 26; i++){
    
        int v = sam.ch[u][i];
        if(v) dp[u] += dfs2(v);
    }
    return dp[u];
}

void build(int t){
    
    for(int i = 1; i <= sam.tot; i++){
    
        if(sam.fa[i]) g[sam.fa[i]].emplace_back(i);
    }
    if(!t){
    
        for(int i = 1; i <= sam.tot; i++) dp[i] = siz[i] = 1;
    } else {
    
        dfs1(1);
    }
    dp[1] = siz[1] = 0;
    dfs2(1);
}

void query(int u, int k){
    
    if(k > dp[u]){
    
        cout << -1;
        return;
    }
    if(k <= siz[u]) return;
    k -= siz[u];
    for(int i = 0; i < 26; i++){
    
        int v = sam.ch[u][i];
        if(k > dp[v]) k -= dp[v];
        else{
    
            cout << (char)('a' + i);
            query(v, k);
            return;
        }
    }
}

inline void solve(){
    
    string s; cin >> s;
    for(int i = 0; i < s.size(); i++) sam.extend(s[i] - 'a');
    int t, k; cin >> t >> k;
    build(t);
    query(1, k);
}

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