【洛谷P3975】弦论

题目

题目链接：https://www.luogu.com.cn/problem/P3975
为了提高智商，ZJY 开始学习弦论。这一天，她在《String theory》中看到了这样一道问题：对于一个给定的长度为 (n) 的字符串，求出它的第 (k) 小子串是什么。你能帮帮她吗?
(t) 为 (0) 则表示不同位置的相同子串算作一个，(t) 为 (1) 则表示不同位置的相同子串算作多个。

思路

建出 SAM 之后，我们分开讨论 (t=0) 和 (t=1) 的点。
对于 (t=0) 的点，我们不难发现经过一个点本质不同的子串的数量就是其 SAM 的 DAG 上所能到达的点。因为在这个类之前所有串的后缀都是一致的，而从 DAG 上这个类可以到达的点也就意味着有多少个不同的“前缀”。所以我们直接在 DAG 上记忆化搜索即可。
对于 (t=1) 的点，经过一个节点的位置不同的子串的数量就是他所代表的类的集合大小。那么我们可以在 parent 树上拓扑求出每一个节点的子树大小，然后再记忆化搜索。
询问的时候，枚举当前位每一个字符，如果数量不够就减去它的数量，否则输出它继续递归。
时间复杂度 (O(nk))，其中 (k=26)。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N=1000010;
int n,k,t,a[N];
ll cnt[N];
char s[N];

struct SAM
{
	int tot,last,ch[N][26],fa[N],len[N];
	ll siz[N];
	SAM() { tot=last=1; }
	
	void ins(int c)
	{
		int p=last,np=++tot; 
		last=np; len[np]=len[p]+1; siz[np]=1;
		for (;p && !ch[p][c];p=fa[p]) ch[p][c]=np;
		if (!p) fa[np]=1;
		else
		{
			int q=ch[p][c];
			if (len[q]==len[p]+1) fa[np]=q;
			else
			{
				int nq=++tot;
				len[nq]=len[p]+1; fa[nq]=fa[q];
				for (int i=0;i<26;i++) ch[nq][i]=ch[q][i];
				fa[q]=fa[np]=nq;
				for (;p && ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
			}
		}
	}
	
	void topsort()
	{
		for (int i=1;i<=tot;i++) cnt[len[i]]++;
		for (int i=1;i<=tot;i++) cnt[i]+=cnt[i-1];
		for (int i=1;i<=tot;i++) a[cnt[len[i]]--]=i;
		for (int i=tot;i>=1;i--) siz[fa[a[i]]]+=siz[a[i]];
		memset(cnt,0,sizeof(cnt));
	}
	
	void dfs(int x)
	{
		if (cnt[x]) return;
		if (t) cnt[x]=siz[x];
			else cnt[x]=siz[x]=(x>1);
		for (int i=0;i<26;i++)
			if (ch[x][i])
			{
				dfs(ch[x][i]);
				cnt[x]+=cnt[ch[x][i]];
			}
	}
	
	void query(int x,int k)
	{
		if (k<=siz[x]) return;
		k-=siz[x];
		for (int i=0;i<26;i++)
			if (k>cnt[ch[x][i]]) k-=cnt[ch[x][i]];
			else
			{
				putchar(i+'a');
				query(ch[x][i],k);
				return;
			}
	}
}sam;

int main()
{
	scanf("%s%d%d",s+1,&t,&k);
	n=strlen(s+1);
	for (int i=1;i<=n;i++)
		sam.ins(s[i]-'a');
	if (t) sam.topsort();
	cnt[1]=sam.siz[1]=0;
	sam.dfs(1);
	if (k>cnt[1]) printf("-1");
		else sam.query(1,k);
	return 0;
}